Reduce code duplication and remove some unused files/functions

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:

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}")

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

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()

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
         """