From 28ef2d870824425f5bfb1efb2947e886029f38c6 Mon Sep 17 00:00:00 2001
From: Itay Mizeretz <itay.mizeretz@guardicore.com>
Date: Thu, 22 Feb 2018 15:43:51 +0200
Subject: [PATCH] return 401 on invalid token

---
 monkey_island/cc/auth.py | 26 +++++++++++++++-----------
 1 file changed, 15 insertions(+), 11 deletions(-)

diff --git a/monkey_island/cc/auth.py b/monkey_island/cc/auth.py
index 510a741ad..7ee787a75 100644
--- a/monkey_island/cc/auth.py
+++ b/monkey_island/cc/auth.py
@@ -1,7 +1,7 @@
 from functools import wraps
 
-import flask_jwt
-from flask_jwt import JWT
+from flask import current_app, abort
+from flask_jwt import JWT, _jwt_required, JWTError
 from werkzeug.security import safe_str_cmp
 
 from cc.island_config import AUTH_ENABLED
@@ -43,12 +43,16 @@ def init_jwt(app):
 
 
 def jwt_required(realm=None):
-    if AUTH_ENABLED:
-        return flask_jwt.jwt_required(realm)
-    else:
-        def wrapper(fn):
-            @wraps(fn)
-            def decorator(*args, **kwargs):
-                return fn(*args, **kwargs)
-            return decorator
-        return wrapper
+    def wrapper(fn):
+        @wraps(fn)
+        def decorator(*args, **kwargs):
+            if AUTH_ENABLED:
+                try:
+                    _jwt_required(realm or current_app.config['JWT_DEFAULT_REALM'])
+                except JWTError:
+                    abort(401)
+            return fn(*args, **kwargs)
+
+        return decorator
+
+    return wrapper