Skip to content

HTTPS clone URL

Subversion checkout URL

You can clone with HTTPS or Subversion.

Download ZIP
Newer
Older
100644 242 lines (187 sloc) 5.922 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
27
28 class BadClient(Exception):
29 pass
30
31 class FeedUnpack(object):
32 def __init__(self):
33 self.buf = bytearray()
34 def __iter__(self):
35 return self
36 def next(self):
37 return self.unpack()
38 def feed(self, data):
39 self.buf.extend(data)
40 def unpack(self):
41 if len(self.buf) < 5:
42 raise StopIteration('No message.')
43
44 ml, opcode = struct.unpack('!iB', buffer(self.buf,0,5))
45 if len(self.buf) < ml:
46 raise StopIteration('No message.')
47
48 data = bytearray(buffer(self.buf, 5, ml-5))
49 del self.buf[:ml]
50 return opcode, data
51
52
53 class FeedConn(EventGen):
54 def __init__(self, conn, addr, db):
55 EventGen.__init__(self)
56 self.conn = conn
57 self.addr = addr
58 self.db = db
59 self.pubchans = set()
60 self.subchans = set()
61 self.idents = set()
62 self.delay = False
63
64 self.rand = struct.pack('<I', random.randint(2**31,2**32-1))
65 self.fu = FeedUnpack()
66
67 conn._on('read', self.io_in)
68 conn._on('close', self.closed)
69
70 self.sendinfo()
71
72 def sendinfo(self):
73 self.conn.write(self.msginfo())
74
75 def auth(self, ident, hash):
76 p = self.db.query('hpfeeds.authkey', {'identifier': str(ident)}, limit=1)
77 p._when(self.checkauth, hash)
78
79 def dbexc(e):
80 logging.critical('Database query exception. {0}'.format(e))
81 self.error('Database query exception.')
82
83 p._except(dbexc)
84
85 self.delay = True
86
87 def checkauth(self, r, hash):
88 if len(r) > 0:
89 akobj = r[0]
90 akhash = hashlib.sha1('{0}{1}'.format(self.rand, akobj['secret'])).hexdigest()
91 if akhash == hash:
92 self.pubchans.update(akobj['publish'])
93 self.subchans.update(akobj['subscribe'])
94 self.idents.add(akobj['identifier'])
95 logging.info('Auth success by {0}.'.format(akobj['identifier']))
96 else:
97 self.error('authfail.')
98 logging.info('Auth failure by {0}.'.format(akobj['identifier']))
99 else:
100 self.error('authfail.')
101
102 self.delay = False
103 self.io_in(b'')
104
105 def closed(self, reason):
106 logging.debug('Connection closed, {0}'.format(reason))
107 self._event('close', self)
108
109 def io_in(self, data):
110 self.fu.feed(data)
111 if self.delay:
112 return
113 try:
114 for opcode, data in self.fu:
115 if opcode == OP_PUBLISH:
116 rest = buffer(data, 0)
117 ident, rest = rest[1:1+ord(rest[0])], buffer(rest, 1+ord(rest[0]))
118 chan, rest = rest[1:1+ord(rest[0])], buffer(rest, 1+ord(rest[0]))
119
120 if not ident in self.idents:
121 self.error('identfail.')
122 continue
123
124 if not chan in self.pubchans:
125 self.error('accessfail.')
126 continue
127
128 self._event('publish', self, chan, data)
129 elif opcode == OP_SUBSCRIBE:
130 rest = buffer(data, 0)
131 ident, chan = rest[1:1+ord(rest[0])], rest[1+ord(rest[0]):]
132
133 if not ident in self.idents:
134 self.error('identfail.')
135 continue
136
137 if not chan in self.subchans:
138 self.error('accessfail.')
139 continue
140
141 self._event('subscribe', self, chan)
142 elif opcode == OP_AUTH:
143 rest = buffer(data, 0)
144 ident, hash = rest[1:1+ord(rest[0])], rest[1+ord(rest[0]):]
145 self.auth(ident, hash)
146
147 except BadClient:
148 self.conn.close()
149 logging.warn('Disconnecting bad client: {0}'.format(self.addr))
150
151 def forward(self, data):
152 self.conn.write(self.msghdr(OP_PUBLISH, data))
153
154 def error(self, emsg):
155 self.conn.write(self.msgerror(emsg))
156
157 def msgerror(self, emsg):
158 return self.msghdr(OP_ERROR, emsg)
159
160 def msginfo(self):
161 return self.msghdr(OP_INFO, '{0}{1}{2}'.format(chr(len(FBNAME)%0xff), FBNAME, self.rand))
162
163 def msghdr(self, op, data):
164 return struct.pack('!iB', 5+len(data), op) + data
165
166
167 class FeedBroker(object):
168 def __init__(self):
169 self.ready = False
170
171 self.db = None
172 self.db = MongoConn(MONGOIP, MONGOPORT)
173 self.db._on('ready', self._dbready)
174 self.db._on('close', self._dbclose)
175
176 self.listener = listenplain(host=FBIP, port=FBPORT)
177 self.listener._on('close', self._lclose)
178 self.listener._on('connection', self._newconn)
179
180 self.listener2 = listenplain(host=FBIP, port=FBPORT+1)
181 self.listener2._on('close', self._lclose)
182 self.listener2._on('connection', self._newconnplain)
183
184 self.connections = set()
185 self.subscribermap = collections.defaultdict(list)
186 self.conn2chans = collections.defaultdict(list)
187
188 def _dbready(self):
189 self.ready = True
190 logging.info('Database ready.')
191
192 def _dbclose(self, e):
193 logging.critical('Database connection closed ({0}). Exiting.'.format(e))
194 unloop()
195
196 def _lclose(self, e):
197 logging.critical('Listener closed ({0}). Exiting.'.format(e))
198 unloop()
199
200 def _newconn(self, c, addr):
201 logging.debug('Connection from {0}.'.format(addr))
202 fc = FeedConn(c, addr, self.db)
203 self.connections.add(fc)
204 fc._on('close', self._connclose)
205 fc._on('subscribe', self._subscribe)
206 fc._on('publish', self._publish)
207
208 def _newconnplain(self, c, addr):
209 logging.debug('Connection from {0}.'.format(addr))
210 fc = FeedConnPlain(c, addr, self.db)
211 self.connections.add(fc)
212 fc._on('close', self._connclose)
213 fc._on('subscribe', self._subscribe)
214 fc._on('publish', self._publish)
215
216 def _connclose(self, c):
217 self.connections.remove(c)
218 for chan in self.conn2chans[c]:
219 self.subscribermap[chan].remove(c)
220
221 def _publish(self, c, chan, data):
222 logging.debug('broker publish to {0} by {1}'.format(chan, c.addr))
223 for c2 in self.subscribermap[chan]:
224 if c2 == c:
225 continue
226 c2.forward(data)
227
228 def _subscribe(self, c, chan):
229 logging.debug('broker subscribe to {0} by {1}'.format(chan, c.addr))
230 self.subscribermap[chan].append(c)
231 self.conn2chans[c].append(chan)
232
233 def main():
234 fb = FeedBroker()
235
236 loop()
237 return 0
238
239 if __name__ == '__main__':
240 sys.exit(main())
241
Something went wrong with that request. Please try again.