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