From 4bb569dd897cb54e4d23390d307b0d863827296b Mon Sep 17 00:00:00 2001
From: Itay Mizeretz <itay.mizeretz@guardicore.com>
Date: Thu, 22 Feb 2018 10:33:37 +0200
Subject: [PATCH] Add JWT authentication to backend

---
 monkey_island/cc/app.py                       | 20 +++++--
 monkey_island/cc/auth.py                      | 54 +++++++++++++++++++
 monkey_island/cc/island_config.py             |  4 ++
 .../monkey_island_pip_requirements.txt        |  1 +
 monkey_island/requirements.txt                |  1 +
 5 files changed, 75 insertions(+), 5 deletions(-)
 create mode 100644 monkey_island/cc/auth.py

diff --git a/monkey_island/cc/app.py b/monkey_island/cc/app.py
index 9c85f6230..a5601ba74 100644
--- a/monkey_island/cc/app.py
+++ b/monkey_island/cc/app.py
@@ -1,22 +1,26 @@
+import os
 from datetime import datetime
+
 import bson
-from bson.json_util import dumps
-from flask import Flask, send_from_directory, redirect, make_response
 import flask_restful
+from bson.json_util import dumps
+from flask import Flask, send_from_directory, make_response
 from werkzeug.exceptions import NotFound
 
+from cc.auth import init_jwt
 from cc.database import mongo
+from cc.island_config import AUTH_EXPIRATION_TIME
 from cc.resources.client_run import ClientRun
-from cc.resources.monkey import Monkey
+from cc.resources.edge import Edge
 from cc.resources.local_run import LocalRun
-from cc.resources.telemetry import Telemetry
+from cc.resources.monkey import Monkey
 from cc.resources.monkey_configuration import MonkeyConfiguration
 from cc.resources.monkey_download import MonkeyDownload
 from cc.resources.netmap import NetMap
-from cc.resources.edge import Edge
 from cc.resources.node import Node
 from cc.resources.report import Report
 from cc.resources.root import Root
+from cc.resources.telemetry import Telemetry
 from cc.resources.telemetry_feed import TelemetryFeed
 from cc.services.config import ConfigService
 
@@ -70,6 +74,12 @@ def init_app(mongo_url):
     api.representations = {'application/json': output_json}
 
     app.config['MONGO_URI'] = mongo_url
+
+    app.config['SECRET_KEY'] = os.urandom(32)
+    app.config['JWT_AUTH_URL_RULE'] = '/api/auth'
+    app.config['JWT_EXPIRATION_DELTA'] = AUTH_EXPIRATION_TIME
+
+    init_jwt(app)
     mongo.init_app(app)
 
     with app.app_context():
diff --git a/monkey_island/cc/auth.py b/monkey_island/cc/auth.py
new file mode 100644
index 000000000..510a741ad
--- /dev/null
+++ b/monkey_island/cc/auth.py
@@ -0,0 +1,54 @@
+from functools import wraps
+
+import flask_jwt
+from flask_jwt import JWT
+from werkzeug.security import safe_str_cmp
+
+from cc.island_config import AUTH_ENABLED
+
+__author__ = 'itay.mizeretz'
+
+
+class User(object):
+    def __init__(self, id, username, password):
+        self.id = id
+        self.username = username
+        self.password = password
+
+    def __str__(self):
+        return "User(id='%s')" % self.id
+
+
+users = [
+    User(1, 'monkey', 'infection')
+]
+username_table = {u.username: u for u in users}
+userid_table = {u.id: u for u in users}
+
+
+def authenticate(username, password):
+    user = username_table.get(username, None)
+    if user and safe_str_cmp(user.password.encode('utf-8'), password.encode('utf-8')):
+        return user
+
+
+def identity(payload):
+    user_id = payload['identity']
+    return userid_table.get(user_id, None)
+
+
+def init_jwt(app):
+    if AUTH_ENABLED:
+        JWT(app, authenticate, identity)
+
+
+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
diff --git a/monkey_island/cc/island_config.py b/monkey_island/cc/island_config.py
index 0a8f33bac..87fe97263 100644
--- a/monkey_island/cc/island_config.py
+++ b/monkey_island/cc/island_config.py
@@ -1,5 +1,9 @@
+from datetime import timedelta
+
 __author__ = 'itay.mizeretz'
 
 ISLAND_PORT = 5000
 DEFAULT_MONGO_URL = "mongodb://localhost:27017/monkeyisland"
 DEBUG_SERVER = False
+AUTH_ENABLED = True
+AUTH_EXPIRATION_TIME = timedelta(hours=1)
diff --git a/monkey_island/deb-package/monkey_island_pip_requirements.txt b/monkey_island/deb-package/monkey_island_pip_requirements.txt
index 404aad8b0..4b4e9d523 100644
--- a/monkey_island/deb-package/monkey_island_pip_requirements.txt
+++ b/monkey_island/deb-package/monkey_island_pip_requirements.txt
@@ -8,6 +8,7 @@ click
 flask
 Flask-Pymongo
 Flask-Restful
+Flask-JWT
 jsonschema
 netifaces
 ipaddress
diff --git a/monkey_island/requirements.txt b/monkey_island/requirements.txt
index 9d8bfbfb8..6aea32b84 100644
--- a/monkey_island/requirements.txt
+++ b/monkey_island/requirements.txt
@@ -8,6 +8,7 @@ click
 flask
 Flask-Pymongo
 Flask-Restful
+Flask-JWT
 jsonschema
 netifaces
 ipaddress