├── .gitignore
├── Makefile
├── README.md
├── conf.py
├── flask_app.py
├── message_proxy.py
├── nginx.conf
├── requirements.txt
├── static
    └── js
    │   ├── cookies.js
    │   ├── user.js
    │   └── websocket.js
├── supervisord.conf
├── templates
    └── index.html
└── websocket_server.py


/.gitignore:
--------------------------------------------------------------------------------
1 | *.pyc
2 | log
3 | *~
4 | \#*
5 | run
6 | __pycache__
7 | 


--------------------------------------------------------------------------------
/Makefile:
--------------------------------------------------------------------------------
 1 | .PHONY: all paths supervisord
 2 | 
 3 | SHELL=/bin/bash
 4 | 
 5 | # Supervisord only runs under Python 2, but that is no problem, even
 6 | # if everything else runs under Python 3.
 7 | SUPERVISORD=supervisord
 8 | 
 9 | all: supervisord
10 | 
11 | paths:
12 | 	mkdir -p log run
13 | 	mkdir -p log/sv_child
14 | 
15 | supervisord: paths
16 | 	$(SUPERVISORD) -c supervisord.conf
17 | 
18 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # Message Flow Demo
 2 | 
 3 | Demonstrates how to set up a Python web service with WebSockets for
 4 | pushing messages from server to front-end.
 5 | 
 6 | ## Installation
 7 | 
 8 | - Install nginx
 9 | - Install supervisord
10 | - Edit the Makefile to point to supervisord
11 | - ``pip install -r requirements.txt``
12 | 
13 | ## Running
14 | 
15 | - ``make``
16 | - Connect to ``http://localhost:5000``
17 | 
18 | ## System outline
19 | 
20 | <a href="http://www.gliffy.com/go/publish/9190819">
21 | <img src="http://www.gliffy.com/go/publish/image/9190819/L.png">
22 | </a>
23 | 


--------------------------------------------------------------------------------
/conf.py:
--------------------------------------------------------------------------------
1 | # Secret used to authenticate websockets.  This can be anything you want.
2 | secret = 'kljasdlkahsdlkajhsdlkajhsd71826381623123'
3 | 


--------------------------------------------------------------------------------
/flask_app.py:
--------------------------------------------------------------------------------
 1 | from flask import Flask, render_template, request, jsonify, escape
 2 | 
 3 | import json
 4 | import threading
 5 | import hashlib
 6 | import uuid
 7 | import collections
 8 | import conf
 9 | import jwt
10 | 
11 | 
12 | app = Flask(__name__)
13 | app.secret_key = conf.secret
14 | 
15 | 
16 | import zmq
17 | ctx = zmq.Context()
18 | pub = ctx.socket(zmq.PUB)
19 | pub.connect('ipc:///tmp/message_flow_in')
20 | 
21 | 
22 | # Task IDs
23 | class TID:
24 |     OK = 'OK'
25 |     ERROR = 'ERROR'
26 |     DONE = 'TASK DONE'
27 | 
28 | 
29 | def push(user, task_id, data):
30 |     """Push message to `user` over websocket.
31 | 
32 |     """
33 |     pub.send(b"0 " + json.dumps({'username': user,
34 |                                  'id': task_id,
35 |                                  'data': data}).encode('utf-8'))
36 | 
37 | 
38 | def long_task(user, message):
39 |     """This is an example of a task that takes long to execute.
40 | 
41 |     """
42 |     import time
43 |     try:
44 |         from time import monotonic
45 |     except ImportError:
46 |         from time import time as monotonic
47 |     import random
48 | 
49 |     m0 = monotonic()
50 |     time.sleep(2 + random.random())
51 |     m1 = monotonic()
52 | 
53 |     processed_message = 'Message processed, "{}" in {:.2f} seconds'.format(escape(message), m1 - m0)
54 |     push(user, TID.DONE, processed_message)
55 | 
56 | 
57 | @app.route('/send', methods=['POST'])
58 | def send():
59 |     message = request.form.get('message', '')
60 | 
61 |     if not message:
62 |         return jsonify(status=TID.ERROR, data='No message submitted')
63 |     else:
64 |         threading.Thread(target=long_task, args=(get_username(), message)).start()
65 |         return jsonify(status=TID.OK,
66 |                        data='Message submitted for processing')
67 | 
68 | 
69 | # Modify this for your specific application
70 | #
71 | # !! Certainly don't get usernames from the client like I do here--that would be highly
72 | # insecure
73 | def get_username():
74 |     return request.cookies.get('username')
75 | 
76 | 
77 | # !!!
78 | # This API call should **only be callable by logged in users**
79 | # !!!
80 | @app.route('/socket_auth_token', methods=['GET'])
81 | def socket_auth_token():
82 |     return jwt.encode({'username': get_username()}, app.secret_key)
83 | 
84 | 
85 | @app.route('/')
86 | def index():
87 |     return render_template('index.html', username=get_username())
88 | 


--------------------------------------------------------------------------------
/message_proxy.py:
--------------------------------------------------------------------------------
 1 | # http://zguide.zeromq.org/page:all#The-Dynamic-Discovery-Problem
 2 | 
 3 | import zmq
 4 | 
 5 | IN = 'ipc:///tmp/message_flow_in'
 6 | OUT = 'ipc:///tmp/message_flow_out'
 7 | 
 8 | context = zmq.Context()
 9 | 
10 | feed_in = context.socket(zmq.XSUB)
11 | feed_in.bind(IN)
12 | 
13 | feed_out = context.socket(zmq.XPUB)
14 | feed_out.bind(OUT)
15 | 
16 | print('[message_proxy] Forwarding messages between {} and {}'.format(IN, OUT))
17 | zmq.proxy(feed_in, feed_out)
18 | 


--------------------------------------------------------------------------------
/nginx.conf:
--------------------------------------------------------------------------------
 1 | http {
 2 |   upstream websocket_server {
 3 |     server localhost:4567;
 4 |   }
 5 | 
 6 |   server {
 7 |     listen 5000;
 8 | 
 9 |     location / {
10 |       proxy_pass http://unix:run/flask_app.sock:/;
11 | 
12 |       proxy_set_header        Host $http_host;
13 |       proxy_set_header        X-Real-IP $remote_addr;
14 |       proxy_set_header        X-Forwarded-For $proxy_add_x_forwarded_for;
15 |       proxy_set_header        X-Forwarded-Proto $scheme;
16 |     }
17 | 
18 |     location /websocket {
19 |         proxy_pass http://websocket_server/websocket;
20 |         proxy_http_version 1.1;
21 |         proxy_set_header Upgrade $http_upgrade;
22 |         proxy_set_header Connection "upgrade";
23 |         proxy_read_timeout 60s;
24 |     }
25 | 
26 |     error_log log/error.log;
27 |     access_log log/nginx-access.log;
28 | 
29 |   }
30 | 
31 | }
32 | 
33 | events {
34 | }
35 | 
36 | pid run/nginx.pid;
37 | error_log log/error.log;
38 | 


--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
1 | tornado>=3.0
2 | supervisor
3 | pyzmq
4 | pyjwt
5 | 


--------------------------------------------------------------------------------
/static/js/cookies.js:
--------------------------------------------------------------------------------
 1 | // From http://www.quirksmode.org/js/cookies.html
 2 | 
 3 | function createCookie(name, value, minutes) {
 4 |     var expires = "";
 5 |     if (minutes) {
 6 |         var date = new Date();
 7 |         date.setTime(date.getTime() + (minutes * 60 * 1000));
 8 |         expires = "; expires=" + date.toGMTString();
 9 |     }
10 |     document.cookie = name + "=" + value + expires + "; path=/";
11 | }
12 | 
13 | function readCookie(name) {
14 |     var nameEQ = name + "=";
15 |     var ca = document.cookie.split(';');
16 |     for(var i=0; i < ca.length; i++) {
17 |         var c = ca[i];
18 |         while (c.charAt(0) == ' ') c = c.substring(1, c.length);
19 |         if (c.indexOf(nameEQ) == 0) return c.substring(nameEQ.length, c.length);
20 |     }
21 |     return null;
22 | }
23 | 
24 | function eraseCookie(name) {
25 |     createCookie(name, "", -1);
26 | }
27 | 


--------------------------------------------------------------------------------
/static/js/user.js:
--------------------------------------------------------------------------------
1 | function generateUserCookie() {
2 |     var username = "test" +
3 |         Math.floor(Math.random() * 9000) + 1000 + "@myservice.org";
4 |     createCookie("username", username);
5 | 
6 |     return username;
7 | }
8 | 


--------------------------------------------------------------------------------
/static/js/websocket.js:
--------------------------------------------------------------------------------
 1 | function getTime() {
 2 |   var date = new Date();
 3 |   var n = date.toDateString();
 4 |   return date.toLocaleTimeString();
 5 | };
 6 | 
 7 | function webSocketAuth() {
 8 |     var auth_token = new Promise(
 9 |         function(resolve, reject) {
10 |             // First, try and read the authentication token from a cookie
11 |             var cookie_token = readCookie('auth_token');
12 |             var no_token = "no_auth_token_user bad_token";
13 | 
14 |             if (cookie_token) {
15 |                 resolve(cookie_token);
16 |             } else {
17 |                 // If not found, ask the server for a new one
18 |                 $.ajax({
19 |                     url: $SCRIPT_ROOT + "/socket_auth_token",
20 |                     statusCode: {
21 |                         // If we get a gateway error, it probably means nginx is being restarted.
22 |                         // Not much we can do, other than wait a bit and continue with a
23 |                         // fake token.
24 |                         405: function() {
25 |                             setTimeout(function() { resolve(no_token); }, 1000);
26 |                         }
27 |                     }})
28 |                     .done(function(data) {
29 |                         createCookie('auth_token', data);
30 |                         resolve(data);
31 |                     })
32 |                     .fail(function() {
33 |                         // Same situation as with 405.
34 |                         setTimeout(function() { resolve(no_token); }, 1000);
35 |                     });
36 |             }
37 |         }
38 |     );
39 | 
40 |     auth_token.then(
41 |         function(token) {
42 |             ws.send(token);
43 |         });
44 | };
45 | 
46 | 
47 | function createSocket(url, messageHandler) {
48 |     var ws = new ReconnectingWebSocket(url);
49 | 
50 |     ws.onopen = function (event) {
51 |         $("#connected").empty().append("Online");
52 |         $("#authenticated").empty().append("No");
53 |     };
54 | 
55 |     ws.onmessage = function (event) {
56 |         var message = event.data;
57 | 
58 |         // Ignore heartbeat signals
59 |         if (message === '<3') {
60 |             return;
61 |         }
62 | 
63 |         var data = JSON.parse(message);
64 |         var id = data["id"];
65 | 
66 |         switch (id) {
67 |         case "AUTH REQUEST":
68 |             webSocketAuth();
69 |             break;
70 |         case "AUTH FAILED":
71 |             eraseCookie('auth_token');
72 |             break;
73 |         case "AUTH OK":
74 |             $("#connected").empty().append("Online");
75 |             $("#authenticated").empty().append("Yes");
76 |             break;
77 |         default:
78 |             messageHandler(data);
79 |         }
80 |     };
81 | 
82 |     ws.onclose = function (event) {
83 |         $("#connected").empty().append("Offline");
84 |         $("#authenticated").empty().append("No");
85 |     };
86 | 
87 |     return ws;
88 | };
89 | 


--------------------------------------------------------------------------------
/supervisord.conf:
--------------------------------------------------------------------------------
 1 | [supervisord]
 2 | logfile=log/supervisord.log
 3 | pidfile=run/supervisord.pid
 4 | nodaemon=true
 5 | childlogdir=log/sv_child
 6 | 
 7 | [program:waitress]
 8 | command=waitress-serve --unix-socket=run/flask_app.sock flask_app:app
 9 | environment=PYTHONUNBUFFERED=1
10 | stdout_logfile=log/waitress.log
11 | redirect_stderr=true
12 | 
13 | [program:nginx]
14 | command=nginx -c nginx.conf -p . -g "daemon off;"
15 | 
16 | [program:message_proxy]
17 | command=/usr/bin/env python message_proxy.py
18 | environment=PYTHONUNBUFFERED=1
19 | stdout_logfile=log/message_proxy.log
20 | redirect_stderr=true
21 | 
22 | [program:websocket]
23 | command=/usr/bin/env python websocket_server.py
24 | environment=PYTHONUNBUFFERED=1
25 | stdout_logfile=log/websocket_server.log
26 | redirect_stderr=true
27 | 
28 | 


--------------------------------------------------------------------------------
/templates/index.html:
--------------------------------------------------------------------------------
  1 | <!doctype html>
  2 | <title>Index</title>
  3 | 
  4 | <head>
  5 |   <script src="http://ajax.googleapis.com/ajax/libs/jquery/1.7/jquery.js"></script>
  6 |   <script src="https://cdnjs.cloudflare.com/ajax/libs/reconnecting-websocket/1.0.0/reconnecting-websocket.min.js"></script>
  7 | 
  8 |   <script src="{{ url_for('static', filename='js/cookies.js')}}"></script>
  9 |   <script src="{{ url_for('static', filename='js/user.js')}}"></script>
 10 |   <script src="{{ url_for('static', filename='js/websocket.js')}}"></script>
 11 | 
 12 |   <script type="text/javascript">
 13 |     var $SCRIPT_ROOT = {{ request.script_root|tojson|safe }};
 14 |   </script>
 15 | 
 16 | </head>
 17 | 
 18 | <h1>WebSocket Demo</h1>
 19 | 
 20 | <ul>
 21 |   <li><b>User: </b><span id="username"></span></li>
 22 |   <li><b>WebSocket Connection: </b><span id="connected">Offline</span></li>
 23 |   <li><b>WebSocket Authenticated: </b><span id="authenticated">No</span></li>
 24 | </ul>
 25 | 
 26 | <form id="messageForm" action="/send" method="post">
 27 |   Message:
 28 |   <input type="text" name="message"/>
 29 | </form>
 30 | 
 31 | <p>API Result:</p>
 32 | <hr/>
 33 | <ol id="result">
 34 | </ol>
 35 | <hr/>
 36 | 
 37 | <p>Incoming:</p>
 38 | <hr/>
 39 | <ol id="incoming">
 40 | </ol>
 41 | <hr/>
 42 | 
 43 | <script type="text/javascript">
 44 |   /* Here we define how incoming events (announced on the WebSocket)
 45 |   are handled */
 46 | 
 47 |   function messageHandler(data) {
 48 |     switch (data["id"]) {
 49 |       case "TASK DONE":
 50 |         $("#incoming")
 51 |           .append("<li>" + getTime() + " [" + data["id"] + "] " + data["data"] + "</li>");
 52 |         break;
 53 |       default:
 54 |         $("#incoming")
 55 |          .append("<li>Unknown message (" + JSON.stringify(data) + ")</li>");
 56 |       }
 57 |   };
 58 | 
 59 | 
 60 |   $(document).ready(function() {
 61 |     /* Fake username insertion.  Remove this when using in the real
 62 |     world, and get the username from the Flask context. */
 63 |     $USER = readCookie("username") || generateUserCookie();
 64 | 
 65 |     /* Set up the WebSocket, and tell it which function will handle
 66 |        incoming messages */
 67 | 
 68 |     // Note: you should be running SSL and change this to wss
 69 |     ws = createSocket("ws://" + {{ request.host|tojson|safe }} + "/websocket",
 70 |                       messageHandler
 71 |     );
 72 | 
 73 | 
 74 |     $("#username").empty().append($USER);
 75 | 
 76 |     /* Setup up the submission form behavior */
 77 | 
 78 |     $("#messageForm").submit(function(event) {
 79 |       event.preventDefault();
 80 | 
 81 |       $.post($SCRIPT_ROOT + "/send",
 82 |              {message: $('input[name="message"]').val()})
 83 | 
 84 |              .success(function(result) {
 85 |                $("#result").append('<li><b>' + getTime() + '</b>' + ' [' +
 86 |                                    result['status'] + '] ' + result['data'] + '</li>');
 87 |              })
 88 | 
 89 |              .error(function(e) {
 90 |                $("#result").append('<li><b>' + e.statusText + '</b></li>');
 91 |              } );
 92 | 
 93 |       return false;
 94 |     });
 95 |   });
 96 | </script>
 97 | 
 98 | 
 99 | </html>
100 | 


--------------------------------------------------------------------------------
/websocket_server.py:
--------------------------------------------------------------------------------
  1 | # encoding: utf-8
  2 | 
  3 | from tornado import websocket, web, ioloop
  4 | import json
  5 | import zmq
  6 | import jwt
  7 | 
  8 | from conf import secret
  9 | 
 10 | ctx = zmq.Context()
 11 | 
 12 | 
 13 | # Could also use: http://aaugustin.github.io/websockets/
 14 | 
 15 | 
 16 | class WebSocket(websocket.WebSocketHandler):
 17 |     participants = set()
 18 | 
 19 |     def __init__(self, *args, **kwargs):
 20 |         websocket.WebSocketHandler.__init__(self, *args, **kwargs)
 21 | 
 22 |         self.authenticated = False
 23 |         self.auth_failures = 0
 24 |         self.username = None
 25 | 
 26 |     def check_origin(self, origin):
 27 |         return True
 28 | 
 29 |     def open(self):
 30 |         if self not in self.participants:
 31 |             self.participants.add(self)
 32 |             self.request_auth()
 33 | 
 34 |     def on_close(self):
 35 |         if self in self.participants:
 36 |             self.participants.remove(self)
 37 | 
 38 |     def on_message(self, auth_token):
 39 |         self.authenticate(auth_token)
 40 |         if not self.authenticated and self.auth_failures < 3:
 41 |             self.request_auth()
 42 | 
 43 |     def request_auth(self):
 44 |         self.auth_failures += 1
 45 |         self.send_json(id="AUTH REQUEST")
 46 | 
 47 |     def send_json(self, **kwargs):
 48 |         self.write_message(json.dumps(kwargs))
 49 | 
 50 |     def authenticate(self, auth_token):
 51 |         try:
 52 |             token_payload = jwt.decode(auth_token, secret)
 53 |             self.username = token_payload['username']
 54 |             self.authenticated = True
 55 |             self.send_json(id='AUTH OK')
 56 |         except DecodeError:
 57 |             self.send_json(id='AUTH FAILED')
 58 | 
 59 |     @classmethod
 60 |     def heartbeat(cls):
 61 |         for p in cls.participants:
 62 |             p.write_message(b'<3')
 63 | 
 64 |     # http://mrjoes.github.io/2013/06/21/python-realtime.html
 65 |     @classmethod
 66 |     def broadcast(cls, data):
 67 |         channel, data = data[0].decode('utf-8').split(" ", 1)
 68 |         user = json.loads(data)["username"]
 69 | 
 70 |         for p in cls.participants:
 71 |             if p.authenticated and p.username == user:
 72 |                 p.write_message(data)
 73 | 
 74 | 
 75 | if __name__ == "__main__":
 76 |     PORT = 4567
 77 |     LOCAL_OUTPUT = 'ipc:///tmp/message_flow_out'
 78 | 
 79 |     import zmq
 80 | 
 81 |     # https://zeromq.github.io/pyzmq/eventloop.html
 82 |     from zmq.eventloop import ioloop, zmqstream
 83 | 
 84 |     ioloop.install()
 85 | 
 86 |     sub = ctx.socket(zmq.SUB)
 87 |     sub.connect(LOCAL_OUTPUT)
 88 |     sub.setsockopt(zmq.SUBSCRIBE, b'')
 89 | 
 90 |     print('[websocket_server] Broadcasting {} to all websockets'.format(LOCAL_OUTPUT))
 91 |     stream = zmqstream.ZMQStream(sub)
 92 |     stream.on_recv(WebSocket.broadcast)
 93 | 
 94 |     server = web.Application([
 95 |         (r'/websocket', WebSocket),
 96 |     ])
 97 |     server.listen(PORT)
 98 | 
 99 |     # We send a heartbeat every 45 seconds to make sure that nginx
100 |     # proxy does not time out and close the connection
101 |     ioloop.PeriodicCallback(WebSocket.heartbeat, 45000).start()
102 | 
103 |     print('[websocket_server] Listening for incoming websocket connections on port {}'.format(PORT))
104 |     ioloop.IOLoop.instance().start()
105 | 


--------------------------------------------------------------------------------