Replace the default function decorators with new socketio specific ones

app/corpora/events.py

@@ -1,9 +1,10 @@
 from app import db, logger, socketio
+from app.decorators import socketio_login_required
 from app.events import connected_sessions
 from app.models import Corpus, User
 from .cqi import CQiClient
 from flask import current_app, request
-from flask_login import current_user, login_required
+from flask_login import current_user
 import math
 
 
@@ -23,13 +24,13 @@ corpus_analysis_clients_status = {}
 
 
 @socketio.on('send_analysis_status_cli')
-@login_required
+@socketio_login_required
 def update_status(response):
     update_analysis_status(response['status'])
 
 
 @socketio.on('corpus_analysis_init')
-@login_required
+@socketio_login_required
 def init_corpus_analysis(corpus_id):
     logger.warning('Initiating corpus analysis.')
     socketio.start_background_task(corpus_analysis_session_handler,
@@ -39,7 +40,7 @@ def init_corpus_analysis(corpus_id):
 
 
 @socketio.on('corpus_analysis_query')
-@login_required
+@socketio_login_required
 def corpus_analysis_query(query):
     update_analysis_status('running')
     logger.warning('Recieved a query.')
@@ -113,7 +114,7 @@ def get_analysis_status():
 
 
 @socketio.on('inspect_match')
-@login_required
+@socketio_login_required
 def inspect_match(message):
     client = corpus_analysis_clients.get(request.sid)
     if client is None:

app/decorators.py

@@ -1,19 +1,34 @@
 from functools import wraps
 from flask import abort
 from flask_login import current_user
+from flask_socketio import disconnect
 from .models import Permission
 
 
-def permission_required(permission):
-    def decorator(f):
-        @wraps(f)
-        def decorated_function(*args, **kwargs):
-            if not current_user.can(permission):
-                abort(403)
-            return f(*args, **kwargs)
-        return decorated_function
-    return decorator
-
-
 def admin_required(f):
-    return permission_required(Permission.ADMIN)(f)
+    @wraps(f)
+    def wrapped(*args, **kwargs):
+        if not current_user.can(Permission.ADMIN):
+            abort(403)
+        return f(*args, **kwargs)
+    return wrapped
+
+
+def socketio_login_required(f):
+    @wraps(f)
+    def wrapped(*args, **kwargs):
+        if not current_user.is_authenticated:
+            disconnect()
+        else:
+            return f(*args, **kwargs)
+    return wrapped
+
+
+def socketio_admin_required(f):
+    @wraps(f)
+    def wrapped(*args, **kwargs):
+        if not current_user.can(Permission.ADMIN):
+            disconnect()
+        else:
+            return f(*args, **kwargs)
+    return wrapped

app/events.py

@@ -1,7 +1,7 @@
 from flask import current_app, request
-from flask_login import current_user, login_required
+from flask_login import current_user
 from . import socketio
-from .decorators import admin_required
+from .decorators import socketio_admin_required, socketio_login_required
 from .models import User
 import json
 import jsonpatch
@@ -16,7 +16,6 @@ connected_sessions = []
 
 
 @socketio.on('connect')
-@login_required
 def connect():
     '''
     ' The Socket.IO module creates a session id (sid) for each request.
@@ -26,7 +25,6 @@ def connect():
 
 
 @socketio.on('disconnect')
-@login_required
 def disconnect():
     '''
     ' On disconnect the session id gets removed from the connected sessions
@@ -36,7 +34,7 @@ def disconnect():
 
 
 @socketio.on('user_ressources_init')
-@login_required
+@socketio_login_required
 def subscribe_user_ressources():
     socketio.start_background_task(user_ressource_session_handler,
                                    current_app._get_current_object(),
@@ -44,8 +42,8 @@ def subscribe_user_ressources():
 
 
 @socketio.on('foreign_user_ressources_init')
-@login_required
+@socketio_login_required
-@admin_required
+@socketio_admin_required
 def subscribe_foreign_user_ressources(user_id):
     socketio.start_background_task(user_ressource_session_handler,
                                    current_app._get_current_object(),