Clean up all the imports.
[quix0rs-apt-p2p.git] / krpc.py
1 ## Copyright 2002-2003 Andrew Loewenstern, All Rights Reserved
2 # see LICENSE.txt for license information
3
4 from bencode import bencode, bdecode
5 from time import asctime
6 import sys
7 from traceback import format_exception
8
9 from twisted.internet.defer import Deferred
10 from twisted.internet import protocol
11 from twisted.internet import reactor
12
13 KRPC_TIMEOUT = 20
14
15 KRPC_ERROR = 1
16 KRPC_ERROR_METHOD_UNKNOWN = 2
17 KRPC_ERROR_RECEIVED_UNKNOWN = 3
18 KRPC_ERROR_TIMEOUT = 4
19
20 # commands
21 TID = 't'
22 REQ = 'q'
23 RSP = 'r'
24 TYP = 'y'
25 ARG = 'a'
26 ERR = 'e'
27
28 class hostbroker(protocol.DatagramProtocol):       
29     def __init__(self, server):
30         self.server = server
31         # this should be changed to storage that drops old entries
32         self.connections = {}
33         
34     def datagramReceived(self, datagram, addr):
35         #print `addr`, `datagram`
36         #if addr != self.addr:
37         c = self.connectionForAddr(addr)
38         c.datagramReceived(datagram, addr)
39         #if c.idle():
40         #    del self.connections[addr]
41
42     def connectionForAddr(self, addr):
43         if addr == self.addr:
44             raise Exception
45         if not self.connections.has_key(addr):
46             conn = self.protocol(addr, self.server, self.transport)
47             self.connections[addr] = conn
48         else:
49             conn = self.connections[addr]
50         return conn
51
52     def makeConnection(self, transport):
53         protocol.DatagramProtocol.makeConnection(self, transport)
54         tup = transport.getHost()
55         self.addr = (tup.host, tup.port)
56
57 ## connection
58 class KRPC:
59     noisy = 1
60     def __init__(self, addr, server, transport):
61         self.transport = transport
62         self.factory = server
63         self.addr = addr
64         self.tids = {}
65         self.mtid = 0
66
67     def datagramReceived(self, str, addr):
68         # bdecode
69         try:
70             msg = bdecode(str)
71         except Exception, e:
72             if self.noisy:
73                 print "response decode error: " + `e`
74         else:
75             #if self.noisy:
76             #    print msg
77             # look at msg type
78             if msg[TYP]  == REQ:
79                 ilen = len(str)
80                 # if request
81                 #       tell factory to handle
82                 f = getattr(self.factory ,"krpc_" + msg[REQ], None)
83                 msg[ARG]['_krpc_sender'] =  self.addr
84                 if f and callable(f):
85                     try:
86                         ret = apply(f, (), msg[ARG])
87                     except Exception, e:
88                         ## send error
89                         out = bencode({TID:msg[TID], TYP:ERR, ERR :`format_exception(type(e), e, sys.exc_info()[2])`})
90                         olen = len(out)
91                         self.transport.write(out, addr)
92                     else:
93                         if ret:
94                             #   make response
95                             out = bencode({TID : msg[TID], TYP : RSP, RSP : ret})
96                         else:
97                             out = bencode({TID : msg[TID], TYP : RSP, RSP : {}})
98                         #       send response
99                         olen = len(out)
100                         self.transport.write(out, addr)
101
102                 else:
103                     if self.noisy:
104                         print "don't know about method %s" % msg[REQ]
105                     # unknown method
106                     out = bencode({TID:msg[TID], TYP:ERR, ERR : KRPC_ERROR_METHOD_UNKNOWN})
107                     olen = len(out)
108                     self.transport.write(out, addr)
109                 if self.noisy:
110                     print "%s %s >>> %s - %s %s %s" % (asctime(), addr, self.factory.node.port, 
111                                                     ilen, msg[REQ], olen)
112             elif msg[TYP] == RSP:
113                 # if response
114                 #       lookup tid
115                 if self.tids.has_key(msg[TID]):
116                     df = self.tids[msg[TID]]
117                     #   callback
118                     del(self.tids[msg[TID]])
119                     df.callback({'rsp' : msg[RSP], '_krpc_sender': addr})
120                 else:
121                     print 'timeout ' + `msg[RSP]['id']`
122                     # no tid, this transaction timed out already...
123             elif msg[TYP] == ERR:
124                 # if error
125                 #       lookup tid
126                 if self.tids.has_key(msg[TID]):
127                     df = self.tids[msg[TID]]
128                     #   callback
129                     df.errback(msg[ERR])
130                     del(self.tids[msg[TID]])
131                 else:
132                     # day late and dollar short
133                     pass
134             else:
135                 print "unknown message type " + `msg`
136                 # unknown message type
137                 df = self.tids[msg[TID]]
138                 #       callback
139                 df.errback(KRPC_ERROR_RECEIVED_UNKNOWN)
140                 del(self.tids[msg[TID]])
141                 
142     def sendRequest(self, method, args):
143         # make message
144         # send it
145         msg = {TID : chr(self.mtid), TYP : REQ,  REQ : method, ARG : args}
146         self.mtid = (self.mtid + 1) % 256
147         str = bencode(msg)
148         d = Deferred()
149         self.tids[msg[TID]] = d
150         def timeOut(tids = self.tids, id = msg[TID]):
151             if tids.has_key(id):
152                 df = tids[id]
153                 del(tids[id])
154                 print ">>>>>> KRPC_ERROR_TIMEOUT"
155                 df.errback(KRPC_ERROR_TIMEOUT)
156         reactor.callLater(KRPC_TIMEOUT, timeOut)
157         self.transport.write(str, self.addr)
158         return d
159