Merge branch 'master' of github.com:paradoxxxzero/butterfly

butterfly.server.py

@@ -20,7 +20,11 @@
 import tornado.options
 import tornado.ioloop
 import tornado.httpserver
-import tornado_systemd
+try:
+    from tornado_systemd import SystemdHTTPServer as HTTPServer
+except ImportError:
+    from tornado.httpserver import HTTPServer
+
 import logging
 import webbrowser
 import uuid
@@ -40,6 +44,8 @@ tornado.options.define("unminified", default=False,
 
 tornado.options.define("host", default='localhost', help="Server host")
 tornado.options.define("port", default=57575, type=int, help="Server port")
+tornado.options.define("keepalive_interval", default=30, type=int,
+                        help="Interval between ping packets sent from server to client (in seconds)")
 tornado.options.define("one_shot", default=False,
                        help="Run a one-shot instance. Quit at term close")
 tornado.options.define("shell", help="Shell to execute at login")
@@ -131,6 +137,7 @@ if not os.path.exists(options.ssl_dir):
 def to_abs(file):
     return os.path.join(options.ssl_dir, file)
 
+
 ca, ca_key, cert, cert_key, pkcs12 = map(to_abs, [
     'butterfly_ca.crt', 'butterfly_ca.key',
     'butterfly_%s.crt', 'butterfly_%s.key',
@@ -156,6 +163,7 @@ def read(file):
     with open(file, 'rb') as fd:
         return fd.read()
 
+
 if options.generate_certs:
     from OpenSSL import crypto
     print('Generating certificates for %s (change it with --host)\n' % host)
@@ -298,9 +306,7 @@ else:
 from butterfly import application
 application.butterfly_dir = butterfly_dir
 log.info('Starting server')
-http_server = tornado_systemd.SystemdHTTPServer(
-    application,
-    ssl_options=ssl_opts)
+http_server = HTTPServer(application, ssl_options=ssl_opts)
 http_server.listen(port, address=host)
 
 if http_server.systemd:

butterfly/routes.py

@@ -18,7 +18,9 @@
 
 import json
 import os
+import struct
 import sys
+import time
 from collections import defaultdict
 from mimetypes import guess_type
 from uuid import uuid4
@@ -116,26 +118,47 @@ class ThemeStatic(Route):
         raise tornado.web.HTTPError(404)
 
 
+class KeptAliveWebSocketHandler(tornado.websocket.WebSocketHandler):
+    keepalive_timer = None
+
+    def open(self, *args, **kwargs):
+        self.keepalive_timer = tornado.ioloop.PeriodicCallback(
+            self.send_ping, tornado.options.options.keepalive_interval * 1000)
+
+    def send_ping(self):
+        t = int(time.time())
+        frame = struct.pack('<I', t)  # A ping frame based on time
+        self.log.info("Sending ping frame %s" % t)
+        try:
+            self.ping(frame)
+        except tornado.websocket.WebSocketClosedError:
+            self.keepalive_timer.stop()
+
+    def on_close(self):
+        if self.keepalive_timer is not None:
+            self.keepalive_timer.stop()
+
+
 @url(r'/ctl/session/(?P<session>[^/]+)')
-class TermCtlWebSocket(Route, tornado.websocket.WebSocketHandler):
+class TermCtlWebSocket(Route, KeptAliveWebSocketHandler):
     sessions = defaultdict(list)
     sessions_secure_users = {}
 
     def open(self, session):
+        super(TermCtlWebSocket, self).open(session)
         self.session = session
         self.closed = False
         self.log.info('Websocket /ctl opened %r' % self)
 
     def create_terminal(self):
         socket = utils.Socket(self.ws_connection.stream.socket)
-        opts = tornado.options.options
         user = self.request.query_arguments.get(
             'user', [b''])[0].decode('utf-8')
         path = self.request.query_arguments.get(
             'path', [b''])[0].decode('utf-8')
         secure_user = None
 
-        if not opts.unsecure:
+        if not tornado.options.options.unsecure:
             user = utils.parse_cert(
                 self.ws_connection.stream.socket.getpeercert())
             assert user, 'No user in certificate'
@@ -208,6 +231,7 @@ class TermCtlWebSocket(Route, tornado.websocket.WebSocketHandler):
         self.broadcast(self.session, message, self)
 
     def on_close(self):
+        super(TermCtlWebSocket, self).on_close()
         if self.closed:
             return
         self.closed = True
@@ -215,8 +239,7 @@ class TermCtlWebSocket(Route, tornado.websocket.WebSocketHandler):
         if self in self.sessions[self.session]:
             self.sessions[self.session].remove(self)
 
-        opts = tornado.options.options
-        if opts.one_shot or (
+        if tornado.options.options.one_shot or (
                 self.application.systemd and
                 not sum([
                     len(wsockets)
@@ -225,7 +248,7 @@ class TermCtlWebSocket(Route, tornado.websocket.WebSocketHandler):
 
 
 @url(r'/ws/session/(?P<session>[^/]+)')
-class TermWebSocket(Route, tornado.websocket.WebSocketHandler):
+class TermWebSocket(Route, KeptAliveWebSocketHandler):
     # List of websockets per session
     sessions = defaultdict(list)
 
@@ -236,6 +259,7 @@ class TermWebSocket(Route, tornado.websocket.WebSocketHandler):
     history = {}
 
     def open(self, session):
+        super(TermWebSocket, self).open(session)
         self.set_nodelay(True)
         self.session = session
         self.closed = False
@@ -273,6 +297,7 @@ class TermWebSocket(Route, tornado.websocket.WebSocketHandler):
         Terminal.sessions[self.session].write(message)
 
     def on_close(self):
+        super(TermWebSocket, self).on_close()
         if self.closed:
             return
         self.closed = True

setup.py

@@ -24,8 +24,11 @@ options = dict(
     platforms="Any",
     scripts=['butterfly.server.py', 'scripts/butterfly', 'scripts/b'],
     packages=['butterfly'],
-    install_requires=["tornado>=3.2", "pyOpenSSL", 'tornado_systemd'],
-    extras_requires=["libsass"],
+    install_requires=["tornado>=3.2", "pyOpenSSL"],
+    extras_requires={
+        'themes': ["libsass"],
+        'systemd': ['tornado_systemd']
+    },
     package_data={
         'butterfly': [
             'sass/*.sass',