krpc.py

   1 import airhook
   2 from twisted.internet.defer import Deferred
   3 from twisted.protocols import basic
   4 from bencode import bencode, bdecode
   5 from twisted.internet import reactor
   6 import time
   7
   8 import hash
   9
  10 KRPC_TIMEOUT = 30
  11
  12 KRPC_ERROR = 1
  13 KRPC_ERROR_METHOD_UNKNOWN = 2
  14 KRPC_ERROR_RECEIVED_UNKNOWN = 3
  15 KRPC_ERROR_TIMEOUT = 4
  16
  17 class KRPC(basic.NetstringReceiver):
  18     noisy = 1
  19     def __init__(self):
  20         self.tids = {}
  21
  22     def stringReceived(self, str):
  23         # bdecode
  24         try:
  25             msg = bdecode(str)
  26         except Exception, e:
  27             print "response decode error: " + `e`
  28             self.d.errback()
  29         else:
  30             # look at msg type
  31             if msg['typ']  == 'req':
  32                 ilen = len(str)
  33                 # if request
  34                 #       tell factory to handle
  35                 f = getattr(self.factory ,"krpc_" + msg['req'], None)
  36                 if f and callable(f):
  37                     msg['arg']['_krpc_sender'] =  self.transport.addr
  38                     try:
  39                         ret = apply(f, (), msg['arg'])
  40                     except Exception, e:
  41                         ## send error
  42                         str = bencode({'tid':msg['tid'], 'typ':'err', 'err' :`e`})
  43                         olen = len(str)
  44                         self.sendString(str)
  45                     else:
  46                         if ret:
  47                             #   make response
  48                             str = bencode({'tid' : msg['tid'], 'typ' : 'rsp', 'rsp' : ret})
  49                         else:
  50                             str = bencode({'tid' : msg['tid'], 'typ' : 'rsp', 'rsp' : []})
  51                         #       send response
  52                         olen = len(str)
  53                         self.sendString(str)
  54
  55                 else:
  56                     # unknown method
  57                     str = bencode({'tid':msg['tid'], 'typ':'err', 'err' : KRPC_ERROR_METHOD_UNKNOWN})
  58                     olen = len(str)
  59                     self.sendString(str)
  60                 if self.noisy:
  61                     print "%s %s >>> %s - %s %s %s" % (time.asctime(), self.transport.addr, self.factory.node.port,
  62                                                     ilen, msg['req'], olen)
  63             elif msg['typ'] == 'rsp':
  64                 # if response
  65                 #       lookup tid
  66                 if self.tids.has_key(msg['tid']):
  67                     df = self.tids[msg['tid']]
  68                     #   callback
  69                     df.callback(msg['rsp'])
  70                     del(self.tids[msg['tid']])
  71                 # no tid, perhaps this transaction timed out already...
  72             elif msg['typ'] == 'err':
  73                 # if error
  74                 #       lookup tid
  75                 df = self.tids[msg['tid']]
  76                 #       callback
  77                 df.errback(msg['err'])
  78                 del(self.tids[msg['tid']])
  79             else:
  80                 # unknown message type
  81                 df = self.tids[msg['tid']]
  82                 #       callback
  83                 df.errback(KRPC_ERROR_RECEIVED_UNKNOWN)
  84                 del(self.tids[msg['tid']])
  85
  86     def sendRequest(self, method, args):
  87         # make message
  88         # send it
  89         msg = {'tid' : hash.newID(), 'typ' : 'req',  'req' : method, 'arg' : args}
  90         str = bencode(msg)
  91         self.sendString(str)
  92         d = Deferred()
  93         self.tids[msg['tid']] = d
  94
  95         def timeOut(tids = self.tids, id = msg['tid']):
  96             if tids.has_key(id):
  97                 df = tids[id]
  98                 del(tids[id])
  99                 df.errback(KRPC_ERROR_TIMEOUT)
 100         reactor.callLater(KRPC_TIMEOUT, timeOut)
 101         return d
 102