Skip to content

HTTPS clone URL

Subversion checkout URL

You can clone with HTTPS or Subversion.

Download ZIP
Newer
Older
100644 277 lines (215 sloc) 6.791 kb
2e58a01 @rep initial import
authored
1 #!/usr/bin/env python
2
3 import sys
4
5 import struct
6 import hashlib
7 import collections
8 import random
9
10 import logging
11 logging.basicConfig(level=logging.DEBUG)
12
13 from evnet import loop, unloop, listenplain, EventGen
14 from evnet.mongodb import MongoConn
15
16 FBIP = '0.0.0.0'
17 FBPORT = 10001
18 FBNAME = '@hp1'
19 MONGOIP = '127.0.0.1'
20 MONGOPORT = 27017
21
22 OP_ERROR = 0
23 OP_INFO = 1
24 OP_AUTH = 2
25 OP_PUBLISH = 3
26 OP_SUBSCRIBE = 4
34a014f @malphx Add a new protocol opcode: Unsubscribe + fix a small typo
malphx authored
27 OP_UNSUBSCRIBE = 5
2e58a01 @rep initial import
authored
28
7ba9966 @rep digest instead of hexdigest
authored
29 MAXBUF = 1024**2
30 SIZES = {
31 OP_ERROR: 5+MAXBUF,
32 OP_INFO: 5+256+20,
33 OP_AUTH: 5+256+20,
34 OP_PUBLISH: 5+MAXBUF,
35 OP_SUBSCRIBE: 5+256*2,
34a014f @malphx Add a new protocol opcode: Unsubscribe + fix a small typo
malphx authored
36 OP_UNSUBSCRIBE: 5+256*2,
7ba9966 @rep digest instead of hexdigest
authored
37 }
38
39
2e58a01 @rep initial import
authored
40 class BadClient(Exception):
41 pass
42
43 class FeedUnpack(object):
44 def __init__(self):
45 self.buf = bytearray()
46 def __iter__(self):
47 return self
48 def next(self):
49 return self.unpack()
50 def feed(self, data):
51 self.buf.extend(data)
52 def unpack(self):
53 if len(self.buf) < 5:
54 raise StopIteration('No message.')
55
56 ml, opcode = struct.unpack('!iB', buffer(self.buf,0,5))
7ba9966 @rep digest instead of hexdigest
authored
57 if ml > SIZES.get(opcode, MAXBUF):
58 raise BadClient('Not respecting MAXBUF.')
59
2e58a01 @rep initial import
authored
60 if len(self.buf) < ml:
61 raise StopIteration('No message.')
62
63 data = bytearray(buffer(self.buf, 5, ml-5))
64 del self.buf[:ml]
65 return opcode, data
66
67
68 class FeedConn(EventGen):
69 def __init__(self, conn, addr, db):
70 EventGen.__init__(self)
71 self.conn = conn
72 self.addr = addr
73 self.db = db
74 self.pubchans = set()
75 self.subchans = set()
76 self.idents = set()
77 self.delay = False
78
79 self.rand = struct.pack('<I', random.randint(2**31,2**32-1))
80 self.fu = FeedUnpack()
81
82 conn._on('read', self.io_in)
83 conn._on('close', self.closed)
84
85 self.sendinfo()
86
87 def sendinfo(self):
88 self.conn.write(self.msginfo())
89
90 def auth(self, ident, hash):
34a014f @malphx Add a new protocol opcode: Unsubscribe + fix a small typo
malphx authored
91 p = self.db.query('hpfeeds.auth_key', {'identifier': str(ident)}, limit=1)
2e58a01 @rep initial import
authored
92 p._when(self.checkauth, hash)
93
94 def dbexc(e):
95 logging.critical('Database query exception. {0}'.format(e))
96 self.error('Database query exception.')
97
98 p._except(dbexc)
99
100 self.delay = True
101
102 def checkauth(self, r, hash):
103 if len(r) > 0:
104 akobj = r[0]
7ba9966 @rep digest instead of hexdigest
authored
105 akhash = hashlib.sha1('{0}{1}'.format(self.rand, akobj['secret'])).digest()
2e58a01 @rep initial import
authored
106 if akhash == hash:
107 self.pubchans.update(akobj['publish'])
108 self.subchans.update(akobj['subscribe'])
109 self.idents.add(akobj['identifier'])
110 logging.info('Auth success by {0}.'.format(akobj['identifier']))
111 else:
112 self.error('authfail.')
113 logging.info('Auth failure by {0}.'.format(akobj['identifier']))
114 else:
115 self.error('authfail.')
116
117 self.delay = False
118 self.io_in(b'')
119
120 def closed(self, reason):
121 logging.debug('Connection closed, {0}'.format(reason))
122 self._event('close', self)
123
124 def io_in(self, data):
125 self.fu.feed(data)
126 if self.delay:
127 return
128 try:
129 for opcode, data in self.fu:
130 if opcode == OP_PUBLISH:
131 rest = buffer(data, 0)
132 ident, rest = rest[1:1+ord(rest[0])], buffer(rest, 1+ord(rest[0]))
133 chan, rest = rest[1:1+ord(rest[0])], buffer(rest, 1+ord(rest[0]))
134
135 if not ident in self.idents:
136 self.error('identfail.')
137 continue
138
139 if not chan in self.pubchans:
140 self.error('accessfail.')
141 continue
142
143 self._event('publish', self, chan, data)
144 elif opcode == OP_SUBSCRIBE:
145 rest = buffer(data, 0)
146 ident, chan = rest[1:1+ord(rest[0])], rest[1+ord(rest[0]):]
147
148 if not ident in self.idents:
149 self.error('identfail.')
150 continue
151
152 if not chan in self.subchans:
153 self.error('accessfail.')
154 continue
155
156 self._event('subscribe', self, chan)
34a014f @malphx Add a new protocol opcode: Unsubscribe + fix a small typo
malphx authored
157 elif opcode == OP_UNSUBSCRIBE:
158 rest = buffer(data, 0)
159 ident, chan = rest[1:1+ord(rest[0])], rest[1+ord(rest[0]):]
160
161 if not ident in self.idents:
162 self.error('identfail.')
163 continue
164
165 if not chan in self.subchans:
166 self.error('accessfail.')
167 continue
168
169 self._event('unsubscribe', self, chan)
2e58a01 @rep initial import
authored
170 elif opcode == OP_AUTH:
171 rest = buffer(data, 0)
172 ident, hash = rest[1:1+ord(rest[0])], rest[1+ord(rest[0]):]
173 self.auth(ident, hash)
174
175 except BadClient:
176 self.conn.close()
177 logging.warn('Disconnecting bad client: {0}'.format(self.addr))
178
179 def forward(self, data):
180 self.conn.write(self.msghdr(OP_PUBLISH, data))
181
182 def error(self, emsg):
183 self.conn.write(self.msgerror(emsg))
184
185 def msgerror(self, emsg):
186 return self.msghdr(OP_ERROR, emsg)
187
188 def msginfo(self):
189 return self.msghdr(OP_INFO, '{0}{1}{2}'.format(chr(len(FBNAME)%0xff), FBNAME, self.rand))
190
191 def msghdr(self, op, data):
192 return struct.pack('!iB', 5+len(data), op) + data
193
194
195 class FeedBroker(object):
196 def __init__(self):
197 self.ready = False
198
199 self.db = None
200 self.db = MongoConn(MONGOIP, MONGOPORT)
201 self.db._on('ready', self._dbready)
202 self.db._on('close', self._dbclose)
203
204 self.listener = listenplain(host=FBIP, port=FBPORT)
205 self.listener._on('close', self._lclose)
206 self.listener._on('connection', self._newconn)
207
208 self.listener2 = listenplain(host=FBIP, port=FBPORT+1)
209 self.listener2._on('close', self._lclose)
210 self.listener2._on('connection', self._newconnplain)
211
212 self.connections = set()
213 self.subscribermap = collections.defaultdict(list)
214 self.conn2chans = collections.defaultdict(list)
215
216 def _dbready(self):
217 self.ready = True
218 logging.info('Database ready.')
219
220 def _dbclose(self, e):
221 logging.critical('Database connection closed ({0}). Exiting.'.format(e))
222 unloop()
223
224 def _lclose(self, e):
225 logging.critical('Listener closed ({0}). Exiting.'.format(e))
226 unloop()
227
228 def _newconn(self, c, addr):
229 logging.debug('Connection from {0}.'.format(addr))
230 fc = FeedConn(c, addr, self.db)
231 self.connections.add(fc)
232 fc._on('close', self._connclose)
233 fc._on('subscribe', self._subscribe)
34a014f @malphx Add a new protocol opcode: Unsubscribe + fix a small typo
malphx authored
234 fc._on('unsubscribe', self._unsubscribe)
2e58a01 @rep initial import
authored
235 fc._on('publish', self._publish)
236
237 def _newconnplain(self, c, addr):
238 logging.debug('Connection from {0}.'.format(addr))
239 fc = FeedConnPlain(c, addr, self.db)
240 self.connections.add(fc)
241 fc._on('close', self._connclose)
242 fc._on('subscribe', self._subscribe)
34a014f @malphx Add a new protocol opcode: Unsubscribe + fix a small typo
malphx authored
243 fc._on('unsubscribe', self._unsubscribe)
2e58a01 @rep initial import
authored
244 fc._on('publish', self._publish)
245
246 def _connclose(self, c):
247 self.connections.remove(c)
248 for chan in self.conn2chans[c]:
249 self.subscribermap[chan].remove(c)
250
251 def _publish(self, c, chan, data):
252 logging.debug('broker publish to {0} by {1}'.format(chan, c.addr))
253 for c2 in self.subscribermap[chan]:
254 if c2 == c:
255 continue
256 c2.forward(data)
257
258 def _subscribe(self, c, chan):
259 logging.debug('broker subscribe to {0} by {1}'.format(chan, c.addr))
260 self.subscribermap[chan].append(c)
261 self.conn2chans[c].append(chan)
34a014f @malphx Add a new protocol opcode: Unsubscribe + fix a small typo
malphx authored
262
263 def _unsubscribe(self, c, chan):
264 logging.debug('broker unsubscribe to {0} by {1}'.format(chan, c.addr))
265 self.subscribermap[chan].remove(c)
266 self.conn2chans[c].remove(chan)
2e58a01 @rep initial import
authored
267
268 def main():
269 fb = FeedBroker()
270
271 loop()
272 return 0
273
274 if __name__ == '__main__':
275 sys.exit(main())
276
Something went wrong with that request. Please try again.