improve websocket startup

burpui/app.py

@@ -41,6 +41,8 @@ def create_app(conf=None, verbose=0, logfile=None, **kwargs):
                    - cli (bool): Are we running the CLI. Default is False.
                    - reverse_proxy (bool): Are we behind a reverse-proxy.
                    Default is True if gunicorn is True
+                   - websocket_server (bool): Are we running the websocket
+                   server. Default is False
     :type kwargs: dict
 
     :returns: A :class:`burpui.server.BUIServer` object
@@ -381,7 +383,8 @@ def create_app(conf=None, verbose=0, logfile=None, **kwargs):
     app.login_manager.init_app(app)
 
     # Create WebSocket server
-    create_websocket(app, websocket_server, celery_worker, gunicorn, cli)
+    if create_websocket(app, websocket_server, celery_worker, gunicorn, cli):
+        return app
 
     # Create celery app if enabled
     create_celery(app, warn=False)
@@ -464,8 +467,6 @@ def create_app(conf=None, verbose=0, logfile=None, **kwargs):
                 session_manager.delete_session()
         return response
 
-    return app
-
     if app.demo and SENTRY_AVAILABLE:
         @app.errorhandler(500)
         def internal_server_error(error):
@@ -476,5 +477,7 @@ def create_app(conf=None, verbose=0, logfile=None, **kwargs):
                 public_dsn=sentry.client.get_public_dsn('https')
             )
 
+    return app
+
 
 init = create_app

burpui/cli.py

@@ -43,7 +43,7 @@ if VERBOSE:
 
 # UNITTEST is used to skip the burp-2 requirements for modes != server
 UNITTEST = os.getenv('BUI_MODE') not in ['server', 'manage', 'celery', 'legacy', 'ws']
-CLI = os.getenv('BUI_MODE') not in ['server', 'legacy', 'ws']
+CLI = os.getenv('BUI_MODE') not in ['server', 'legacy']
 
 try:
     app = create_app(

burpui/extensions.py

@@ -215,7 +215,7 @@ def create_websocket(myapp, websocket_server=False, celery_worker=False,
     :type myapp: :class:`burpui.server.BUIServer`
     """
     if cli and not websocket_server:
-        return
+        return False
     broker = myapp.ws_broker
     if broker is not False:
         if not broker or broker is True:
@@ -251,7 +251,7 @@ def create_websocket(myapp, websocket_server=False, celery_worker=False,
     # myapp.config['WS_ASYNC_MODE'] = 'threading' if not gunicorn else None
 
     if celery_worker:
-        return
+        return False
 
     # if you are not a celery worker, we can patch the flask server
     try:
@@ -269,6 +269,9 @@ def create_websocket(myapp, websocket_server=False, celery_worker=False,
     if myapp.config['WITH_WS'] or websocket_server:
         from .ws.namespace import BUINamespace
         socketio.on_namespace(BUINamespace('/ws'))
+        return True
+
+    return False
 
 
 def create_celery(myapp, warn=True):