[Backend] Created access decorators for all user types

flask-backend/api/helpers/access_decorators.py

@@ -0,0 +1,51 @@
+from functools import wraps
+from flask_login import current_user
+
+# Use this decorator on routes which only admins should be able to access
+def admin_only(fun):
+  @wraps(fun)
+  def wrap(*args, **kwargs):
+    print(current_user)
+    # If user is not logged in
+    if not current_user.is_authenticated:
+      return 'Unauthorized. Please log in as Admin', 401
+
+    # If user is not admin
+    if current_user.role != 'admin':
+      return 'Unauthorized. Only Admins can access this route', 401
+
+    # Else continue
+    return fun(*args, **kwargs)
+  return wrap
+
+# Use this decorator on routes which only extractors should be able to access
+def extractor_only(fun):
+  @wraps(fun)
+  def wrap(*args, **kwargs):
+    # If user is not logged in
+    if not current_user.is_authenticated:
+      return 'Unauthorized. Please log in as Extractor', 401
+
+    # If user is not admin
+    if current_user.role != 'extractor':
+      return 'Unauthorized. Only Extractors can access this route', 401
+
+    # Else continue
+    return fun(*args, **kwargs)
+  return wrap
+
+# Use this decorator on routes which only management should be able to access
+def management_only(fun):
+  @wraps(fun)
+  def wrap(*args, **kwargs):
+    # If user is not logged in
+    if not current_user.is_authenticated:
+      return 'Unauthorized. Please log in as Management', 401
+
+    # If user is not admin
+    if current_user.role != 'management':
+      return 'Unauthorized. Only Management can access this route', 401
+
+    # Else continue
+    return fun(*args, **kwargs)
+  return wrap

flask-backend/api/routes/case.py

@@ -4,6 +4,7 @@
 from flask_login import login_required, current_user
 from ..models.models import Case, CaseSchema
 from .. import db
+from ..helpers.access_decorators import management_only
 
 ROOT_DIR = os.getcwd()
 
@@ -13,17 +14,20 @@
 case = Blueprint('case', __name__, url_prefix='/case')
 
 @case.route('/count', methods=["GET"])
+@management_only
 def count():
     return jsonify({'status':200,
                     'total_users':Case.query.count()})
 
 @case.route('/list', methods=["GET"])
+@management_only
 def list():
     all_cases = Case.query.order_by(Case.timestamp).all()
     result = cases_schema.dump(all_cases)
     return jsonify(result)
 
 @case.route('/delete', methods=['POST'])
+@management_only
 def deletecase():
     # check if case_name is provided
     try:
@@ -42,6 +46,7 @@ def deletecase():
     return 'case deleted', 202
 
 @case.route('/open/<case_name>', methods=["GET"])
+@management_only
 def openCase(case_name):
     os.chdir('../../..')
     path = os.getcwd()+'/data/'+case_name
@@ -50,6 +55,7 @@ def openCase(case_name):
     return files
 
 @case.route('/list-files/<case_name>/<folder_name>', methods=["GET"])
+@management_only
 def openFolder(case_name, folder_name):
     os.chdir('../../..')
     path = os.getcwd()+'/data/'+case_name+'/'+folder_name
@@ -58,6 +64,7 @@ def openFolder(case_name, folder_name):
     return files
 
 @case.route('/list-files/<case_name>/<folder_name>/<file_name>', methods=["GET"])
+@management_only
 def openFile(case_name, folder_name, file_name):
     os.chdir('../../..')
     File = os.getcwd()+'/data/'+case_name+'/'+folder_name+'/'+file_name

flask-backend/api/routes/extraction.py

@@ -8,6 +8,7 @@
 from flask_login import login_required, current_user
 from ..models.models import Case, CaseSchema
 from .. import db
+from ..helpers.access_decorators import extractor_only
 ROOT_DIR = os.getcwd()
 
 case_schema = CaseSchema()
@@ -18,6 +19,7 @@
 
 
 @extraction.route('/list_devices', methods=["GET"])
+@extractor_only
 def list_devices():
     with open(os.devnull, 'wb') as devnull:
         subprocess.check_call([adb_path, 'start-server'], stdout=devnull,
@@ -43,6 +45,7 @@ def list_devices():
     return json.dumps(devices)
 
 @extraction.route('/extract_data', methods=["POST"])
+@extractor_only
 def extract():
 
     # if no data is provided at all

flask-backend/api/routes/user.py

@@ -3,6 +3,7 @@
 from flask import Blueprint, render_template, jsonify, request
 from flask_login import login_required, current_user
 from ..models.models import User, UserSchema
+from ..helpers.access_decorators import admin_only
 from werkzeug.security import generate_password_hash, check_password_hash
 from .. import db
 from sqlalchemy import update
@@ -88,7 +89,7 @@ def create_user(): # Add only admin can create functionality, once deployed on a
         email = str(req['email'])
         password = str(req['password'])
         name = str(req['name'])
-        role = ''.join(sorted(str(req['role'])))
+        role = ''.join(str(req['role']))
     except KeyError as err:
         return f'please provide {str(err)}', 400
 
@@ -109,81 +110,77 @@ def create_user(): # Add only admin can create functionality, once deployed on a
 
 # Route for admin to add user 
 @user.route('/add-user', methods=['POST'])
+@admin_only
 @login_required
 def add_users():
-    if(current_user.has_admin == False): 
-        try:
-            req = request.get_json()
-            email = str(req['email'])
-            password = str(req['password'])
-            name = str(req['name'])
-            role = ''.join(sorted(str(req['role'])))
-            timestamp = int(time.time())
-        except:
-            return 'Please provide all parameters', 409
-        user = User.query.filter_by(email=email).first()
-
-        if user:
-            return 'Email address already exists', 409
-        elif role == 'adimn':
-            return 'You cannot create a user of role admin', 409
-
-        new_user = User(email=email, name=name, password=generate_password_hash(password, method='sha256'), role=role, timestamp=timestamp)
-        new_user.admin = current_user.email 
-        db.session.add(new_user)
-        db.session.commit()
+    try:
+        req = request.get_json()
+        email = str(req['email'])
+        password = str(req['password'])
+        name = str(req['name'])
+        role = ''.join(sorted(str(req['role'])))
+        timestamp = int(time.time())
+    except:
+        return 'Please provide all parameters', 409
+    user = User.query.filter_by(email=email).first()
 
-        return 'user created', 202
-    return "You can't add users, you are not an admin", 409
+    if user:
+        return 'Email address already exists', 409
+    elif role == 'adimn':
+        return 'You cannot create a user of role admin', 409
+
+    new_user = User(email=email, name=name, password=generate_password_hash(password, method='sha256'), role=role, timestamp=timestamp)
+    new_user.admin = current_user.email 
+    db.session.add(new_user)
+    db.session.commit()
+
+    return 'user created', 202
 
 
 # Route for admin to view all his users
 @user.route('/all-users', methods=['GET'])
+@admin_only
 @login_required
 def all_users():
-    if current_user.role == 'adimn':
-        all_users = User.query.filter_by(admin=current_user.email).order_by(User.timestamp).all()
-        result = users_schema.dump(all_users)
-        return jsonify(result)
-    return 'You are not admin', 409
+    all_users = User.query.filter_by(admin=current_user.email).order_by(User.timestamp).all()
+    result = users_schema.dump(all_users)
+    return jsonify(result)
 
 # Route for admin to delete a user
 @user.route('/remove-user', methods=['POST'])
+@admin_only
 @login_required
 def remove_user():
-    if current_user.role == 'adimn':
-        try:
-            req = request.get_json()
-            email = str(req['email'])
-        except:
-            return 'Please provide all parameters', 409
-        user = User.query.filter_by(admin=current_user.email).filter_by(email=email).first()
-        if user:
-            db.session.delete(user)
-            db.session.commit()
-            return f"User {user.email} removed."
-        return 'User not found.', 409
-    return 'You are not an admin', 409
+    try:
+        req = request.get_json()
+        email = str(req['email'])
+    except:
+        return 'Please provide all parameters', 409
+    user = User.query.filter_by(admin=current_user.email).filter_by(email=email).first()
+    if user:
+        db.session.delete(user)
+        db.session.commit()
+        return f"User {user.email} removed."
+    return 'User not found.', 409
 
 
 # Route to udate role of an user
 @user.route('/role-update', methods=['POST'])
+@admin_only
 @login_required
 def roleupdate():
-    if current_user.role == 'adimn':
-        try: 
-            req = request.get_json()
-            email = str(req['email'])
-            newrole = ''.join(sorted(str(req['role'])))
-        except:
-            return 'Please provide all parameters', 409
-        user = User.query.filter_by(admin=current_user.email).filter_by(email=email).first()
-        if user:
-            user.role = newrole
-            db.session.commit()
-            return f"User {user.email} has role {newrole} now."
-        return 'User not found.', 409
-    return 'You are not an admin.', 409
+    try: 
+        req = request.get_json()
+        email = str(req['email'])
+        newrole = ''.join(sorted(str(req['role'])))
+    except:
+        return 'Please provide all parameters', 409
+    user = User.query.filter_by(admin=current_user.email).filter_by(email=email).first()
+    if user:
+        user.role = newrole
+        db.session.commit()
+        return f"User {user.email} has role {newrole} now."
+    return 'User not found.', 409
 
 
 @user.route('/delete', methods=['POST'])