/
tornado_server.py
178 lines (149 loc) · 5.96 KB
/
tornado_server.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
from urllib.parse import urlparse
from tornado.ioloop import IOLoop
from tornado.web import Application, RequestHandler, HTTPError
from tornado.options import define, parse_command_line, options
from tornado.websocket import WebSocketHandler, WebSocketClosedError
from django.core.signing import TimestampSigner, BadSignature, SignatureExpired
from tornado.httpserver import HTTPServer
from collections import defaultdict
from redis import Redis
from tornadoredis import Client
from tornadoredis.pubsub import BaseSubscriber
from django.utils.crypto import constant_time_compare
import logging
import signal
import time
import json
import uuid
import os
import hashlib
define('debug', default=False, type=bool, help='Run in debug mode')
define('port', default=8080, type=int, help='Server port')
define('allowed_hosts', default='localhost:8080', multiple=True,
help='Allowed hosts for cross-domain connections')
class UpdateHandler(RequestHandler):
"""Handle updates that come in via Django"""
def post(self, model, pk):
self._broadcast(model, pk, 'add')
def put(self, model, pk):
self._broadcast(model, pk, 'update')
def delete(self, model, pk):
self._broadcast(model, pk, 'remove')
def _broadcast(self, model, pk, action):
signature = self.request.headers.get('X-Signature', None)
if not signature:
raise HTTPError(400)
try:
result = self.application.signer.unsign(signature, max_age=60*1)
except (BadSignature, SignatureExpired):
raise HTTPError(400)
else:
expected = '{method}:{url}:{body}'.format(
method=self.request.method.lower(),
url=self.request.full_url(),
body=hashlib.sha256(self.request.body).hexdigest(),
)
if not constant_time_compare(result, expected):
raise HTTPError(400)
try:
body = json.loads(self.request.body.decode('utf-8'))
except ValueError:
body = None
message = json.dumps({
'model': model,
'id': pk,
'action': action,
'body': body,
})
self.application.broadcast(message)
self.write('OK')
class SprintHandler(WebSocketHandler):
"""Handles real-time updates for the board"""
def check_origin(self, origin):
allowed = super().check_origin(origin)
parsed = urlparse(origin.lower())
matched = any(parsed.netloc == host for host in options.allowed_hosts)
return options.debug or allowed or matched
def open(self):
"""Subscribe to sprint updates on a new connection"""
self.sprint = None
channel = self.get_argument('channel', None)
if not channel:
self.close()
else:
try:
self.sprint = self.application.signer.unsign(
channel, max_age=60 * 30)
except (BadSignature, SignatureExpired):
self.close()
else:
self.uid = uuid.uuid4().hex
self.application.add_subscriber(self.sprint, self)
def on_message(self, message):
"""Broadcast updates to other interested clients"""
if self.sprint is not None:
self.application.broadcast(message, channel=self.sprint, sender=self)
def on_close(self):
"""Remove subscription"""
if self.sprint is not None:
self.application.remove_subscriber(self.sprint, self)
class RedisSubscriber(BaseSubscriber):
def on_message(self, msg):
"""Handle new message on Redis channel"""
if msg and msg.kind == 'message':
try:
message = json.loads(msg.body)
sender = message['sender']
message = message['message']
except (ValueError, KeyError):
message = msg.body
sender = None
subscribers = list(self.subscribers[msg.channel].keys())
for subscriber in subscribers:
if sender is None or sender != subscriber.uid:
try:
subscriber.write_message(msg.body)
except WebSocketClosedError:
# Remove dead peer
self.unsubscribe(msg.channel, subscriber)
super().on_message(msg)
class ScrumApplication(Application):
def __init__(self, **kwargs):
routes = [
(r'/socket', SprintHandler),
(r'/(?P<model>task|sprint|user)/(?P<pk>[0-9]+)', UpdateHandler),
]
super().__init__(routes, **kwargs)
self.subscriber = RedisSubscriber(Client())
self.publisher = Redis()
self._key = os.environ.get('TORNADO_SECRET',
'f56A89be7@37714e0!d890z103b^4f6k380b+25')
self.signer = TimestampSigner(self._key)
def add_subscriber(self, channel, subscriber):
self.subscriber.subscribe(['all', channel], subscriber)
def remove_subscriber(self, channel, subscriber):
self.subscriber.unsubscribe(channel, subscriber)
self.subscriber.unsubscribe('all', subscriber)
def broadcast(self, message, channel=None, sender=None):
channel = 'all' if channel is None else channel
messsage = json.dumps({
'sender': sender and sender.uid,
'message': message
})
self.publisher.publish(channel, message)
def shutdown(server):
ioloop = IOLoop.instance()
logging.info('Stopping server')
server.stop()
def finalize():
ioloop.stop()
logging.info('Stopped')
ioloop.add_timeout(time.time() + 1.5, finalize)
if __name__ == '__main__':
parse_command_line()
application = ScrumApplication(debug=options.debug)
server = HTTPServer(application)
server.listen(options.port)
signal.signal(signal.SIGINT, lambda sig, frame: shutdown(server))
logging.info('Starting server on localhost:{0}'.format(options.port))
IOLoop.instance().start()