From 22ad839bd0458584494bd9c7d5d6e2fc035f66ac Mon Sep 17 00:00:00 2001 From: Carter Brainerd <0xCB@protonmail.com> Date: Sat, 25 Nov 2023 13:55:50 -0500 Subject: [PATCH] Reduce code duplication and remove some unused files/functions --- client/cli/command.py | 1 + client/comms.py | 75 +++++-------- server/config/admin/routes.yaml | 5 + server/create_operator.py | 9 ++ server/maliketh/builder/builder.py | 10 +- server/maliketh/crypto/aes.py | 93 ----------------- server/maliketh/crypto/x509.py | 20 ---- server/maliketh/listeners/admin.py | 162 ++++++++++------------------- server/maliketh/listeners/utils.py | 31 ++++++ server/maliketh/models.py | 9 ++ 10 files changed, 142 insertions(+), 273 deletions(-) delete mode 100644 server/maliketh/crypto/aes.py delete mode 100644 server/maliketh/crypto/x509.py create mode 100644 server/maliketh/listeners/utils.py diff --git a/client/cli/command.py b/client/cli/command.py index 6a5b5f9..115bf8c 100644 --- a/client/cli/command.py +++ b/client/cli/command.py @@ -111,6 +111,7 @@ def handle_builder(args: List[str], config: OperatorConfig) -> None: """ if len(args) < 1: logger.error("Please provide an action and a field") + return if args[0] == "set": if len(args) < 2: diff --git a/client/comms.py b/client/comms.py index 7fbb63d..c02b050 100644 --- a/client/comms.py +++ b/client/comms.py @@ -32,13 +32,7 @@ def check_auth_token(config: OperatorConfig) -> bool: """ Check if the auth token is valid """ - url = f"http://{config.c2}:{config.c2_port}/op/auth/token/status" - headers = { - "Authentication": f"Bearer {config.auth_token}", - } - - response = requests.get(url, headers=headers) - return response.json()["status"] + return send_authenticated_request("GET", "/op/auth/token/status", config).json().get("status") def server_auth(ip: str, port: int, name: str, login_secret: str) -> ServerAuthResponse: @@ -102,19 +96,29 @@ def ensure_token(config: OperatorConfig) -> None: config.auth_token = handle_server_auth(config) +def send_authenticated_request(method: str, endpoint: str, config: OperatorConfig, **request_kwargs) -> requests.Response: + """ + Build and send an authenticated response to the given endpoint. The C2 and authentication information + will be extracted from `config`. + """ + url = f"http://{config.c2}:{config.c2_port}{endpoint}" + headers = { + "Authorization": f"Bearer {config.auth_token}", + } + + return requests.request( + method=method, + url=url, + headers=headers, + **request_kwargs) + def list_implants(config: OperatorConfig) -> list: """ List all the implants """ try: ensure_token(config) - - url = f"http://{config.c2}:{config.c2_port}/op/implant/list" - headers = { - "Authorization": f"Bearer {config.auth_token}", - } - - response = requests.get(url, headers=headers) + response = send_authenticated_request("GET", "/op/implant/list", config) return response.json()["implants"] except Exception as e: logger.error("Failed to list implants") @@ -146,12 +150,7 @@ def get_tasks(config: OperatorConfig) -> List[Dict[Any, Any]]: try: ensure_token(config) - url = f"http://{config.c2}:{config.c2_port}/op/tasks/list" - headers = { - "Authorization": f"Bearer {config.auth_token}", - } - - response = requests.get(url, headers=headers, timeout=120) + response = send_authenticated_request("GET", "/op/tasks/list", config, timeout=120) if response.json()["status"] != True: logger.error("Failed to get tasks") return [] @@ -168,18 +167,13 @@ def add_task( try: ensure_token(config) - url = f"http://{config.c2}:{config.c2_port}/op/tasks/add" - headers = { - "Authorization": f"Bearer {config.auth_token}", - } - data = { "opcode": opcode, "implant_id": implant_id, "args": args, } - response = requests.post(url, headers=headers, json=data) + response = send_authenticated_request("POST", "/op/tasks/add", config, json=data) if response.json()["status"] != True: logger.error("Failed to add task") return {} @@ -194,12 +188,7 @@ def get_task_result(config: OperatorConfig, task_id: str) -> Optional[str]: try: ensure_token(config) - url = f"http://{config.c2}:{config.c2_port}/op/tasks/results/{task_id}" - headers = { - "Authorization": f"Bearer {config.auth_token}", - } - - response = requests.get(url, headers=headers) + response = send_authenticated_request("GET", f"/op/tasks/results/{task_id}", config) if response.json()["status"] != True: logger.error("Failed to get task result") return None @@ -220,12 +209,7 @@ def get_implant_profile(config: OperatorConfig, implant_id: str) -> Dict[str, An ensure_token(config) - url = f"http://{config.c2}:{config.c2_port}/op/implant/config/{implant_id}" - headers = { - "Authorization": f"Bearer {config.auth_token}", - } - - response = requests.get(url, headers=headers) + response = send_authenticated_request("GET", f"/op/implant/config/{implant_id}", config) if response.json()["status"] != True: logger.error("Failed to get implant config") return {} @@ -242,6 +226,7 @@ def update_implant_profile(config: OperatorConfig, implant_id: str, changes: Dic } response = requests.post(url, headers=headers, json=changes) + response = send_authenticated_request("POST", f"/op/implant/config/{implant_id}", config, json=changes) if response.json()["status"] != True: logger.error("Failed to update implant config") return @@ -251,12 +236,7 @@ def kill_implant(config: OperatorConfig, implant_id: str) -> None: ensure_token(config) - url = f"http://{config.c2}:{config.c2_port}/op/implant/kill/{implant_id}" - headers = { - "Authorization": f"Bearer {config.auth_token}", - } - - response = requests.delete(url, headers=headers) + response = send_authenticated_request("DELETE", f"/op/implant/kill/{implant_id}", config) if response.json()["status"] != True: logger.error("Failed to kill implant") return @@ -270,12 +250,7 @@ def build_implant(config: OperatorConfig, build_options: dict) -> str: ensure_token(config) - url = f"http://{config.c2}:{config.c2_port}/op/implant/build" - headers = { - "Authorization": f"Bearer {config.auth_token}", - } - - response = requests.post(url, headers=headers, json=build_options) + response = send_authenticated_request("POST", "/op/implant/build", config, json=build_options) if response.status_code != 200: logger.error("Failed to build implant") logger.error(f"Server response: {response.text}") diff --git a/server/config/admin/routes.yaml b/server/config/admin/routes.yaml index 0c44d88..33feef6 100644 --- a/server/config/admin/routes.yaml +++ b/server/config/admin/routes.yaml @@ -70,3 +70,8 @@ auth_token_status: path: /auth/token/status methods: - GET + +admin_revoke_operator: + path: /admin/revoke_access + methods: + - POST diff --git a/server/create_operator.py b/server/create_operator.py index fb2b691..c13ee8e 100644 --- a/server/create_operator.py +++ b/server/create_operator.py @@ -52,6 +52,14 @@ def main(): ) other = parser.add_option_group("Other") + other.add_option( + "-r", + "--role", + action="store", + dest="role", + default="operator", + help="The role to give to the new user, if applicable. One of \"admin\", \"operator\"." + ) other.add_option( "-v", "--verbose", @@ -114,6 +122,7 @@ def main(): auth_token_expiry=None, created_at=datetime.now().strftime("%Y-%m-%d %H:%M:%S"), rmq_queue=operator_config["rmq_queue"], + role=options.role ) db.session.add(operator) db.session.commit() diff --git a/server/maliketh/builder/builder.py b/server/maliketh/builder/builder.py index f1e3022..1ae9dba 100644 --- a/server/maliketh/builder/builder.py +++ b/server/maliketh/builder/builder.py @@ -1,5 +1,5 @@ from dataclasses import dataclass, asdict, field -from typing import Optional +from typing import Any, Generator, Optional from maliketh.crypto.utils import random_hex import subprocess import os @@ -60,14 +60,14 @@ def operator_name(self): def BuilderOptions(self): return self._builder_options - def operator(self, operator_name: str): + def operator(self, name: str): """ Set the operator name """ - self.operator_name = operator_name + self._operator_name = name return self - def with_options(self, options: BuilderOptions): + def with_options(self, options: "BuilderOptions"): """ Set the build options """ @@ -148,7 +148,7 @@ def build(self) -> Optional[bytes]: return None - def __create_compiler_flags(self) -> str: + def __create_compiler_flags(self) -> Generator[str, Any, Any]: """ Create the string of compiler arguments (-D) to pass to the compiler """ diff --git a/server/maliketh/crypto/aes.py b/server/maliketh/crypto/aes.py deleted file mode 100644 index eab07df..0000000 --- a/server/maliketh/crypto/aes.py +++ /dev/null @@ -1,93 +0,0 @@ -import base64 -import os -from Crypto.Cipher import AES -from cryptography.hazmat.primitives.ciphers.aead import AESGCM - -from maliketh.crypto.utils import pad, random_bytes - - -class GCM: - @staticmethod - def encrypt(data: bytes, aad: bytes, key: bytes) -> bytes: - """ - Encrypt a given byte array using AES-GCM - """ - aesgcm = AESGCM(key) - nonce = os.urandom(12) - return aesgcm.encrypt(nonce, data, aad) + nonce - - @staticmethod - def encrypt_b64(data: bytes, aad: bytes, key: bytes) -> str: - """ - Encrypt a given byte array using AES-GCM and return it as a base64 encoded string - """ - return base64.b64encode(GCM.encrypt(data, aad, key)).decode("utf-8") - - @staticmethod - def decrypt(data: bytes, aad: bytes, key: bytes) -> bytes: - """ - Decrypt a given byte array using AES-GCM - """ - aesgcm = AESGCM(key) - nonce = data[-12:] - data = data[:-12] - return aesgcm.decrypt(nonce, data, aad) - - @staticmethod - def decrypt_b64(data: str, aad: bytes, key: bytes) -> bytes: - """ - Decrypt a given base64 encoded string using AES-GCM - """ - return GCM.decrypt(base64.b64decode(data), aad, key) - - @staticmethod - def gen_key(bit_length=256) -> bytes: - """ - Generate a random AES key - """ - return AESGCM.generate_key(bit_length=bit_length) - - @staticmethod - def gen_key_b64(bit_length=256) -> str: - """ - Generate a random AES key and return it as a base64 encoded string - """ - return base64.b64encode(GCM.gen_key(bit_length)).decode("utf-8") - - @staticmethod - def is_valid_key(key: bytes) -> bool: - try: - AESGCM(key) - return True - except: - return False - - -def generate_aes_key() -> str: - """ - Generate a random AES key and return it as a base64 encoded string - """ - return base64.b64encode(random_bytes(32)).decode("utf-8") - - -def generate_aes_iv() -> str: - """ - Generate a random AES IV and return it as a base64 encoded string - """ - return base64.b64encode(random_bytes(16)).decode("utf-8") - - -def encrypt_aes(data: bytes, key: bytes, iv: bytes) -> bytes: - """ - Encrypt a given byte array using AES - """ - cipher = AES.new(key, AES.MODE_CBC, iv) - return cipher.encrypt(pad(data)) - - -def decrypt_aes(data: bytes, key: bytes, iv: bytes) -> bytes: - """ - Decrypt a given byte array using AES - """ - cipher = AES.new(key, AES.MODE_CBC, iv) - return cipher.decrypt(data) diff --git a/server/maliketh/crypto/x509.py b/server/maliketh/crypto/x509.py deleted file mode 100644 index 2a1dbaa..0000000 --- a/server/maliketh/crypto/x509.py +++ /dev/null @@ -1,20 +0,0 @@ -import os -from typing import Optional - -""" -Generate a CA certificate and private key for use with mutual TLS -Note: We should probably use a crypto library for this, but for now use openssl -""" - - -def generate_ca_cert() -> tuple: - - # check if openssl is installed - if os.system("command -v openssl") != 0: - raise Exception("openssl is not installed") - - ca_key = os.popen("openssl genrsa 2048").read() - ca_cert = os.popen( - "openssl req -new -x509 -nodes -days 365 -key - -subj '/CN=maliketh CA'" - ).read() - return ca_key, ca_cert diff --git a/server/maliketh/listeners/admin.py b/server/maliketh/listeners/admin.py index 5627d68..e2a6fc7 100644 --- a/server/maliketh/listeners/admin.py +++ b/server/maliketh/listeners/admin.py @@ -7,17 +7,18 @@ from maliketh.db import db from maliketh.models import * from maliketh.crypto.ec import * -from maliketh.crypto.utils import random_hex, random_string +from maliketh.crypto.utils import random_hex from maliketh.config import OP_ROUTES from maliketh.opcodes import Opcodes from maliketh.builder.builder import ImplantBuilder, BuilderOptions +from maliketh.listeners.utils import error_json, success_json, create_route from functools import wraps admin = Blueprint("admin", __name__, url_prefix=OP_ROUTES["base_path"]) start_time = datetime.now() -def verify_auth_token(request) -> Optional[str]: +def verify_auth_token(request) -> Optional[Operator]: """ Given a request, verify its Authentication header @@ -51,7 +52,8 @@ def verify_auth_token(request) -> Optional[str]: def verified(func): """ - Decorator to check if the request is authenticated + Decorator to check if the request is authenticated and the operator hasnt had their access revoked. + :param func: The function to decorate :return The decorated function with the operator as a keyword argument """ @@ -61,17 +63,18 @@ def wrapper(*args, **kwargs): # Check if the request is authenticated operator = verify_auth_token(request) if operator is None: - return jsonify({"status": False, "msg": "Not authenticated"}), 401 + return error_json("Not authenticated") + + if operator.revoked: + return error_json("Your access has been revoked. Please contact the admin to resolve.") + kwargs = {**kwargs, "operator": operator} return func(*args, **kwargs) return wrapper -@admin.route( - OP_ROUTES["stats"]["path"], - methods=OP_ROUTES["stats"]["methods"], -) +@create_route(admin, "stats") @verified def server_stats(operator: Operator) -> Any: """ @@ -98,19 +101,15 @@ def server_stats(operator: Operator) -> Any: ) -# @admin.route("/op/auth/token/request", methods=["GET"]) # type: ignore -@admin.route( - OP_ROUTES["request_auth_token"]["path"], - methods=OP_ROUTES["request_auth_token"]["methods"], -) +@create_route(admin, "request_auth_token") def request_token() -> Any: # Check if X-ID and X-Signature header is present if any(x not in request.headers for x in ["X-ID", "X-Signature"]): - return jsonify({"status": False, "message": "Unknown operator key"}), 400 + return error_json("Unknown operator key", 400) # Check if they have content if any(len(request.headers[x]) == 0 for x in ["X-ID", "X-Signature"]): - return jsonify({"status": False, "message": "Unknown operator key"}), 400 + return error_json("Unknown operator key", 400) # X-Signature is in format base64(enc_and_sign(pub_key, operator_signing_key, server_pub_key)) # So we need to decrypt it with the server public key, @@ -119,31 +118,23 @@ def request_token() -> Any: # Get the operator from X-ID operator = Operator.query.filter_by(username=request.headers["X-ID"]).first() if operator is None: - return jsonify({"status": False, "message": "Unknown operator"}), 400 + return error_json("Unknown operator", 400) try: original_message = decrypt_and_verify( bytes(request.headers["X-Signature"], "utf-8"), operator ) except nacl.exceptions.CryptoError: - return jsonify({"status": False, "msg": "Couldn't decrypt signature"}), 400 + return error_json("Couldn't decrypt signature", 400) if original_message is None: - return jsonify({"status": False, "msg": "Couldn't verify signature"}), 400 + return error_json("Couldn't verify signature", 400) original_message = original_message.decode("utf-8") if original_message != operator.login_secret: - return ( - jsonify( - { - "status": False, - "msg": "Unable to verify signature", - } - ), - 400, - ) - + return error_json("Unable to verify signature", 400) + # If we get here, the operator is authenticated # Check if the token is still valid @@ -163,8 +154,6 @@ def request_token() -> Any: ) db.session.commit() - # TODO rabbitmq stuff - return ( jsonify( { @@ -179,11 +168,7 @@ def request_token() -> Any: ) -# @admin.route("/op/auth/token/revoke", methods=["GET"]) # type: ignore -@admin.route( - OP_ROUTES["revoke_auth_token"]["path"], - methods=OP_ROUTES["revoke_auth_token"]["methods"], -) +@create_route(admin, "revoke_auth_token") @verified def revoke_token(operator: Operator) -> Any: operator.auth_token = None # type: ignore @@ -192,22 +177,15 @@ def revoke_token(operator: Operator) -> Any: return jsonify({"status": True}), 200 -@admin.route( - OP_ROUTES["auth_token_status"]["path"], - methods=OP_ROUTES["auth_token_status"]["methods"], -) +@create_route(admin, "auth_token_status") def token_status() -> Any: operator = verify_auth_token(request) if operator is None: - return jsonify({"status": False, "msg": "Not authenticated"}), 401 - return jsonify({"status": True, "msg": "Authenticated"}), 200 + return error_json("Not authenticated") + return success_json("Authenticated") -# @admin.route("/op/tasks/list", methods=["GET"]) # type: ignore -@admin.route( - OP_ROUTES["list_tasks"]["path"], - methods=OP_ROUTES["list_tasks"]["methods"], -) +@create_route(admin, "list_tasks") @verified def list_tasks(operator: Operator) -> Any: """ @@ -217,11 +195,7 @@ def list_tasks(operator: Operator) -> Any: return jsonify({"status": True, "tasks": [x.toJSON() for x in tasks]}), 200 -# @admin.route("/op/tasks/add", methods=["POST"]) # type: ignore -@admin.route( - OP_ROUTES["add_task"]["path"], - methods=OP_ROUTES["add_task"]["methods"], -) +@create_route(admin, "add_task") @verified def add_task(operator: Operator) -> Any: """ @@ -229,7 +203,7 @@ def add_task(operator: Operator) -> Any: """ if request.json is None: - return jsonify({"status": False, "msg": "Invalid request, no JSON body"}), 400 + return error_json("Invalid request, no JSON body", 400) # Get the task task = request.json @@ -238,15 +212,7 @@ def add_task(operator: Operator) -> Any: # Check if the task is valid if any(x not in task for x in required_fields): - return ( - jsonify( - { - "status": False, - "msg": f"Invalid task, missing fields: {', '.join(required_fields - task.keys())}", - } - ), - 400, - ) + return error_json(f"Invalid task, missing fields: {', '.join(required_fields - task.keys())}", 400) # Create the task task = Task.new_task( @@ -259,11 +225,7 @@ def add_task(operator: Operator) -> Any: return jsonify({"status": True, "task": task.toJSON()}), 200 -# @admin.route("/op/tasks/result/", methods=["GET"]) # type: ignore -@admin.route( - OP_ROUTES["task_results"]["path"], - methods=OP_ROUTES["task_results"]["methods"], -) +@create_route(admin, "task_results") @verified def get_task_result(operator: Operator, task_id: str) -> Any: """ @@ -272,19 +234,15 @@ def get_task_result(operator: Operator, task_id: str) -> Any: # Check if this operator owns the task task = Task.query.filter_by(task_id=task_id).first() if task is None: - return jsonify({"status": False, "msg": "Unknown task"}), 400 + return error_json("Unknown task", 400) if task.operator_name != operator.username: - return jsonify({"status": False, "msg": "Unauthorized"}), 401 + return error_json("Unauthorized") return jsonify({"status": True, "result": task.output}), 200 -# @admin.route("/op/tasks/delete/", methods=["DELETE"]) # type: ignore -@admin.route( - OP_ROUTES["delete_task"]["path"], - methods=OP_ROUTES["delete_task"]["methods"], -) +@create_route(admin, "delete_task") @verified def delete_task(operator: Operator, task_id: str) -> Any: """ @@ -293,10 +251,10 @@ def delete_task(operator: Operator, task_id: str) -> Any: # Check if this operator owns the task task = Task.query.filter_by(task_id=task_id).first() if task is None: - return jsonify({"status": False, "msg": "Unknown task"}), 400 + return error_json("Unknown task", 400) if task.operator_name != operator.username: - return jsonify({"status": False, "msg": "Unauthorized"}), 401 + return error_json("Unauthorized") db.session.delete(task) db.session.commit() @@ -304,10 +262,7 @@ def delete_task(operator: Operator, task_id: str) -> Any: return jsonify({"status": True}), 200 -@admin.route( - OP_ROUTES["list_implants"]["path"], - methods=OP_ROUTES["list_implants"]["methods"], -) +@create_route(admin, "list_implants") @verified def list_implants(operator: Operator) -> Any: """ @@ -317,17 +272,14 @@ def list_implants(operator: Operator) -> Any: return jsonify({"status": True, "implants": [x.toJSON() for x in implants]}), 200 -@admin.route( - OP_ROUTES["update_implant_config"]["path"], - methods=OP_ROUTES["update_implant_config"]["methods"], -) +@create_route(admin, "update_implant_config") @verified def update_config(operator: Operator, implant_id: str) -> Any: """ Update the config of an implant. This will trigger a config update task """ if request.json is None: - return jsonify({"status": False, "msg": "Invalid request, no JSON body"}), 400 + return error_json("Invalid request, no JSON body", 400) # Get the task config = request.json @@ -335,7 +287,7 @@ def update_config(operator: Operator, implant_id: str) -> Any: # Update implant's config in the database current_config = ImplantConfig.query.filter_by(implant_id=implant_id).first() if current_config is None: - return jsonify({"status": False, "msg": "Unknown implant"}), 400 + return error_json("Unknown implant", 400) try: @@ -347,10 +299,7 @@ def update_config(operator: Operator, implant_id: str) -> Any: } if len(config) == 0: - return ( - jsonify({"status": False, "msg": "No valid fields found in request"}), - 400, - ) + return error_json("No valid fields found in request", 400) # Update fields present in the request for key, value in config.items(): @@ -363,15 +312,12 @@ def update_config(operator: Operator, implant_id: str) -> Any: ) except Exception as e: print(e) - return jsonify({"status": False, "msg": f"Error updating config: {e}"}), 400 + return error_json(f"Error updating config: {e}", 400) return jsonify({"status": True, "task": task.toJSON()}), 200 -@admin.route( - OP_ROUTES["get_implant_config"]["path"], - methods=OP_ROUTES["get_implant_config"]["methods"], -) +@create_route(admin, "get_implant_config") @verified def get_implant_config(implant_id: str, operator: Operator) -> Any: """ @@ -379,15 +325,12 @@ def get_implant_config(implant_id: str, operator: Operator) -> Any: """ config = ImplantConfig.query.filter_by(implant_id=implant_id).first() if config is None: - return jsonify({"status": False, "msg": "Unknown implant"}), 400 + return error_json("Unknown implant", 400) return jsonify({"status": True, "config": config.toJSON()}), 200 -@admin.route( - OP_ROUTES["kill_implant"]["path"], - methods=OP_ROUTES["kill_implant"]["methods"], -) +@create_route(admin, "kill_implant") @verified def kill_implant(operator: Operator, implant_id: str) -> Any: """ @@ -397,7 +340,7 @@ def kill_implant(operator: Operator, implant_id: str) -> Any: # See if implant exists implant = Implant.query.filter_by(implant_id=implant_id).first() if implant is None: - return jsonify({"status": False, "msg": "Unknown implant"}), 400 + return error_json("Unknown implant", 400) # Create the task Task.new_task( @@ -412,22 +355,31 @@ def kill_implant(operator: Operator, implant_id: str) -> Any: return jsonify({"status": True}), 200 -@admin.route( - OP_ROUTES["build_implant"]["path"], - methods=OP_ROUTES["build_implant"]["methods"], -) +@create_route(admin, "build_implant") @verified def build_implant(operator: Operator) -> Any: """ Build an implant """ if request.json is None: - return jsonify({"status": False, "msg": "Invalid request, no JSON body"}), 400 + return error_json("Invalid request, no JSON body", 400) # Get the build options build_opts_json = request.json builder = ImplantBuilder(operator.username) implant_bytes = builder.with_options(BuilderOptions.from_dict(build_opts_json)).build() + if implant_bytes is None: + return error_json("Error building implant", 500) return jsonify({"status": True, "implant": base64.b64encode(implant_bytes).decode('utf-8')}), 200 + +@create_route(admin, "admin_revoke_operator") +@verified +def admin_revoke_operator(operator: Operator): + """ + Administrative endpoint to revoke operators access. + """ + + if request.json is None: + return error_json("Invalid request, no JSON body", 400) \ No newline at end of file diff --git a/server/maliketh/listeners/utils.py b/server/maliketh/listeners/utils.py new file mode 100644 index 0000000..c2b08ea --- /dev/null +++ b/server/maliketh/listeners/utils.py @@ -0,0 +1,31 @@ +from typing import Callable, Tuple +from flask import jsonify, Response, Blueprint +from functools import wraps +from maliketh.config import OP_ROUTES + + +def basic_status_json(status: bool, msg: str, code:int=401) -> Tuple[Response, int]: + return jsonify({"status": status, "msg": msg}), code + +def error_json(msg: str, code:int=401) -> Tuple[Response, int]: + return basic_status_json(False, msg, code) + +def success_json(msg: str, code:int=200) -> Tuple[Response, int]: + return basic_status_json(True, msg, code) + +def create_route(bp: Blueprint, route_name: str) -> Callable: + """ + Utility function to create a route in a blueprint for a given endpoint. + + :param bp: The flask blueprint to add the route to + :param route_name: The name of the route in the admin config file + """ + def inner(func): + + @bp.route(OP_ROUTES[route_name]["path"], methods=OP_ROUTES[route_name]["methods"]) + @wraps(func) + def wrap(*args, **kwargs): + return func(*args, **kwargs) + + return wrap + return inner diff --git a/server/maliketh/models.py b/server/maliketh/models.py index 69a3d16..b07e363 100644 --- a/server/maliketh/models.py +++ b/server/maliketh/models.py @@ -1,5 +1,6 @@ from dataclasses import dataclass, asdict from datetime import datetime +import enum from typing import Any from maliketh.db import db from maliketh.profile import * @@ -240,6 +241,12 @@ def get_task_by_id(task_id: str): def get_oldest_task_for_implant(implant_id: str): return Task.query.filter_by(implant_id=implant_id, status=CREATED).first() +class OperatorRole(enum.Enum): + ADMIN = 'admin' + OPERATOR = 'operator' + + def __str__(self): + return self.value @dataclass class Operator(db.Model): @@ -256,6 +263,8 @@ class Operator(db.Model): created_at: str = db.Column(db.String) last_login: str = db.Column(db.String) rmq_queue: str = db.Column(db.String) # The RabbitMQ queue name for this operator + role: OperatorRole = db.Column(db.String) + revoked: bool = db.Column(db.Boolean) def toJSON(self): return asdict(self)