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 resetConnection(self):
  23         self.brokenPeer = 0
  24         self._readerState = basic.LENGTH
  25         self._readerLength = 0
  26
  27     def stringReceived(self, str):
  28         # bdecode
  29         try:
  30             msg = bdecode(str)
  31         except Exception, e:
  32             if self.naisy:
  33                 print "response decode error: " + `e`
  34             self.d.errback()
  35         else:
  36             # look at msg type
  37             if msg['typ']  == 'req':
  38                 ilen = len(str)
  39                 # if request
  40                 #       tell factory to handle
  41                 f = getattr(self.factory ,"krpc_" + msg['req'], None)
  42                 if f and callable(f):
  43                     msg['arg']['_krpc_sender'] =  self.transport.addr
  44                     try:
  45                         ret = apply(f, (), msg['arg'])
  46                     except Exception, e:
  47                         ## send error
  48                         str = bencode({'tid':msg['tid'], 'typ':'err', 'err' :`e`})
  49                         olen = len(str)
  50                         self.sendString(str)
  51                     else:
  52                         if ret:
  53                             #   make response
  54                             str = bencode({'tid' : msg['tid'], 'typ' : 'rsp', 'rsp' : ret})
  55                         else:
  56                             str = bencode({'tid' : msg['tid'], 'typ' : 'rsp', 'rsp' : {}})
  57                         #       send response
  58                         olen = len(str)
  59                         self.sendString(str)
  60
  61                 else:
  62                     if self.noisy:
  63                         print "don't know about method %s" % msg['req']
  64                     # unknown method
  65                     str = bencode({'tid':msg['tid'], 'typ':'err', 'err' : KRPC_ERROR_METHOD_UNKNOWN})
  66                     olen = len(str)
  67                     self.sendString(str)
  68                 if self.noisy:
  69                     print "%s %s >>> %s - %s %s %s" % (time.asctime(), self.transport.addr, self.factory.node.port,
  70                                                     ilen, msg['req'], olen)
  71             elif msg['typ'] == 'rsp':
  72                 # if response
  73                 #       lookup tid
  74                 if self.tids.has_key(msg['tid']):
  75                     df = self.tids[msg['tid']]
  76                     #   callback
  77                     del(self.tids[msg['tid']])
  78                     df.callback({'rsp' : msg['rsp'], '_krpc_sender': self.transport.addr})
  79                 # no tid, this transaction timed out already...
  80             elif msg['typ'] == 'err':
  81                 # if error
  82                 #       lookup tid
  83                 df = self.tids[msg['tid']]
  84                 #       callback
  85                 df.errback(msg['err'])
  86                 del(self.tids[msg['tid']])
  87             else:
  88                 # unknown message type
  89                 df = self.tids[msg['tid']]
  90                 #       callback
  91                 df.errback(KRPC_ERROR_RECEIVED_UNKNOWN)
  92                 del(self.tids[msg['tid']])
  93
  94     def sendRequest(self, method, args):
  95         # make message
  96         # send it
  97         msg = {'tid' : hash.newID(), 'typ' : 'req',  'req' : method, 'arg' : args}
  98         str = bencode(msg)
  99         d = Deferred()
 100         self.tids[msg['tid']] = d
 101         def timeOut(tids = self.tids, id = msg['tid']):
 102             if tids.has_key(id):
 103                 df = tids[id]
 104                 del(tids[id])
 105                 print ">>>>>> KRPC_ERROR_TIMEOUT"
 106                 df.errback(KRPC_ERROR_TIMEOUT)
 107         reactor.callLater(KRPC_TIMEOUT, timeOut)
 108         self.sendString(str)
 109         return d
 110