diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 6ba7ec26..7e43bff4 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -20,17 +20,36 @@ repos: - id: mypy additional_dependencies: [pydantic, types-PyYAML, types-requests, types-paramiko, types-tabulate] +- repo: https://github.com/myint/autoflake + rev: 'v2.3.1' + hooks: + - id: autoflake + args: [ + --in-place, + --remove-unused-variables, + --remove-all-unused-imports, + --recursive + ] + - repo: https://github.com/pycqa/flake8 rev: '7.1.0' # pick a git hash / tag to point to hooks: - id: flake8 + args: [ + --max-line-length=110, + ] - repo: https://github.com/hhatto/autopep8 rev: 'v2.3.1' hooks: - id: autopep8 - args: [--max-line-length=110, --diff] - + args: [ + --in-place, + --aggressive, + --aggressive, + --max-line-length=110, + --recursive + ] # PROBLEMS WITH IMPORTS IN PYLINT!!! #- repo: https://github.com/PyCQA/prospector diff --git a/docs/source/developing/command.rst b/docs/source/developing/command.rst index 0b14d4e9..10c84d44 100644 --- a/docs/source/developing/command.rst +++ b/docs/source/developing/command.rst @@ -4,14 +4,14 @@ Adding a New Command ====================== -AttackMate supports extending its functionality by adding new commands. +AttackMate supports extending its functionality by adding new commands. This section details the steps required to integrate a new command. 1. Define the Command Schema ============================= -All Commands in AttackMate inherit from ``BaseCommand``. +All Commands in AttackMate inherit from ``BaseCommand``. To create a new command, define a class in `/src/attackmate/schemas` and register it using the ``@CommandRegistry.register('')`` decorator. For example, to add a ``debug`` command: @@ -21,7 +21,7 @@ For example, to add a ``debug`` command: from typing import Literal from .base import BaseCommand from attackmate.command import CommandRegistry - + @CommandRegistry.register('debug') class DebugCommand(BaseCommand): type: Literal['debug'] @@ -30,7 +30,7 @@ For example, to add a ``debug`` command: wait_for_key: bool = False cmd: str = '' -Registering the command in the ``CommandRegistry`` allows the command to be also instantiated dynamically using the ``Command.create()`` method and is essential to +Registering the command in the ``CommandRegistry`` allows the command to be also instantiated dynamically using the ``Command.create()`` method and is essential to make them usable in external python scripts. @@ -53,16 +53,16 @@ The new command should be handled by an executor in `src/attackmate/executors`` 3. Ensure the Executor Handles the New Command ============================================== - -The ``ExecutorFactory`` class manages and creates executor instances based on command types. + +The ``ExecutorFactory`` class manages and creates executor instances based on command types. It maintains a registry (``_executors``) that maps command type strings to executor classes, allowing for dynamic execution of different command types. -Executors are registered using the ``register_executor`` method, which provides a decorator to associate a command type with a class. +Executors are registered using the ``register_executor`` method, which provides a decorator to associate a command type with a class. When a command is executed, the ``create_executor`` method retrieves the corresponding executor class, filters the constructor arguments based on the class's signature, and then creates an instance. -Accordingly, executors must be registered using the ``@executor_factory.register_executor('')`` decorator. +Accordingly, executors must be registered using the ``@executor_factory.register_executor('')`` decorator. -If the new executor class requires additional initialization arguments, these must be added to the ``_get_executor_config`` method in ``attackmate.py``. -All configurations are always passed to the ``ExecutorFactory``. +If the new executor class requires additional initialization arguments, these must be added to the ``_get_executor_config`` method in ``attackmate.py``. +All configurations are always passed to the ``ExecutorFactory``. The factory filters the provided configurations based on the class constructor signature, ensuring that only the required parameters are used. :: @@ -93,7 +93,7 @@ Update the ``LoopCommand`` schema to include the new command. DebugCommand, # Newly added command # ... other command classes ... ] - + 5. Modify playbook.py to Include the New Command ===================================================== @@ -116,10 +116,3 @@ Once these steps are completed, the new command will be fully integrated into At ===================== Finally, update the documentation in `docs/source/playbook/commands` to include the new command. - - - - - - - diff --git a/docs/source/developing/integration.rst b/docs/source/developing/integration.rst index 514f415b..76b13c23 100644 --- a/docs/source/developing/integration.rst +++ b/docs/source/developing/integration.rst @@ -101,4 +101,4 @@ The return code can be used to determine if the command was successful: if result.returncode == 0: print("Command executed successfully.") else: - print(f"Command failed with return code {result.returncode}") \ No newline at end of file + print(f"Command failed with return code {result.returncode}") diff --git a/docs/source/installation/uv.rst b/docs/source/installation/uv.rst index e3fc212c..99d39b70 100644 --- a/docs/source/installation/uv.rst +++ b/docs/source/installation/uv.rst @@ -59,4 +59,4 @@ Installation Steps .. warning:: Please note that you need to :ref:`sliver-fix` if you want - to use the sliver commands! \ No newline at end of file + to use the sliver commands! diff --git a/docs/source/playbook/commands/vnc.rst b/docs/source/playbook/commands/vnc.rst index 117f384c..1f41adfc 100644 --- a/docs/source/playbook/commands/vnc.rst +++ b/docs/source/playbook/commands/vnc.rst @@ -151,8 +151,3 @@ Execute commands on a remote server via VNC. Uses the `vncdotool bool: def get_user_hash(username: str) -> Optional[str]: """Fetches the hashed password from environment variables.""" - env_var_name = f"USER_{username.upper()}_HASH" + env_var_name = f'USER_{username.upper()}_HASH' return os.getenv(env_var_name) @@ -84,7 +84,7 @@ async def get_current_user(token: str = Depends(api_key_header_scheme)) -> str: token_data = ACTIVE_TOKENS.get(token) if not token_data: - logger.warning(f"Token not found: {token[:5]}...") + logger.warning(f'Token not found: {token[:5]}...') raise credentials_exception username: str = token_data['username'] @@ -99,5 +99,5 @@ async def get_current_user(token: str = Depends(api_key_header_scheme)) -> str: renew_token_expiry(token) - logger.debug(f"Token validated successfully for user: {username}") + logger.debug(f'Token validated successfully for user: {username}') return username diff --git a/remote_rest/client.py b/remote_rest/client.py index c3091041..65f990d0 100644 --- a/remote_rest/client.py +++ b/remote_rest/client.py @@ -36,7 +36,7 @@ def save_token(token: Optional[str]): if token: # This is pretty hacky, client mainly for testing purposes logger.info('updating env var') - logger.info(f"run in your shell: export ATTACKMATE_API_TOKEN={token}") + logger.info(f'run in your shell: export ATTACKMATE_API_TOKEN={token}') else: os.environ.pop(TOKEN_ENV_VAR, None) @@ -67,14 +67,14 @@ def parse_key_value_pairs(items: List[str] | None) -> Dict[str, str]: key, value = item.split('=', 1) result[key.strip()] = value.strip() else: - logging.warning(f"Skipping malformed pair: {item}") + logging.warning(f'Skipping malformed pair: {item}') return result # Login def login(client: httpx.Client, base_url: str, username: str, password: str): """Logs in and saves the token.""" - url = f"{base_url}/login" + url = f'{base_url}/login' logger.info(f"Attempting login for user '{username}' at {url}...") try: # standard form encoding for OAuth2PasswordRequestForm -> expected bei Fastapi @@ -84,46 +84,46 @@ def login(client: httpx.Client, base_url: str, username: str, password: str): token = data.get('access_token') if token: save_token(token) # workaround, export to env var in shell - print(f"Login successful. Token received: {token[:5]}...") + print(f'Login successful. Token received: {token[:5]}...') else: logger.error(' No access token received in response.') sys.exit(1) except httpx.RequestError as e: - logger.error(f"HTTP Request Error during login: {e}") + logger.error(f'HTTP Request Error during login: {e}') sys.exit(1) except httpx.HTTPStatusError as e: - logger.error(f"Login failed: {e.response.status_code}") + logger.error(f'Login failed: {e.response.status_code}') sys.exit(1) except Exception as e: - logger.error(f"Unexpected error during login: {e}", exc_info=True) + logger.error(f'Unexpected error during login: {e}', exc_info=True) sys.exit(1) def get_instance_state_from_server(client: httpx.Client, base_url: str, instance_id: str): """Requests the state of a specific instance.""" - url = f"{base_url}/instances/{instance_id}/state" - logger.info(f"Requesting state for instance {instance_id} at {url}...") + url = f'{base_url}/instances/{instance_id}/state' + logger.info(f'Requesting state for instance {instance_id} at {url}...') try: response = client.get(url, headers=get_auth_headers()) response.raise_for_status() data = response.json() update_token_from_response(data) - print(f"\n State for Instance {instance_id} ") + print(f'\n State for Instance {instance_id} ') print(yaml.dump(data.get('variables', {}), indent=2)) except httpx.RequestError as e: - logger.error(f"HTTP Request Error getting state: {e}") + logger.error(f'HTTP Request Error getting state: {e}') except httpx.HTTPStatusError as e: - logger.error(f"HTTP Status Error getting state: {e.response.status_code} - {e.response.text}") + logger.error(f'HTTP Status Error getting state: {e.response.status_code} - {e.response.text}') except Exception as e: - logger.error(f"Unexpected error getting state: {e}", exc_info=True) + logger.error(f'Unexpected error getting state: {e}', exc_info=True) def run_playbook_yaml( client: httpx.Client, base_url: str, playbook_file: str, debug: bool = False ): """Sends playbook YAML content to the server.""" - url = f"{base_url}/playbooks/execute/yaml" - logger.info(f"Attempting to execute playbook from local file: {playbook_file}") + url = f'{base_url}/playbooks/execute/yaml' + logger.info(f'Attempting to execute playbook from local file: {playbook_file}') try: with open(playbook_file, 'r') as f: playbook_yaml_content = f.read() @@ -147,13 +147,13 @@ def run_playbook_yaml( if not data.get('success'): sys.exit(1) except httpx.RequestError as e: - logger.error(f"HTTP Request Error executing playbook YAML: {e}") + logger.error(f'HTTP Request Error executing playbook YAML: {e}') sys.exit(1) except httpx.HTTPStatusError as e: - logger.error(f"HTTP Status Error (YAML): {e.response.status_code} - {e.response.text}") + logger.error(f'HTTP Status Error (YAML): {e.response.status_code} - {e.response.text}') sys.exit(1) except Exception as e: - logger.error(f"Unexpected error (YAML): {e}", exc_info=True) + logger.error(f'Unexpected error (YAML): {e}', exc_info=True) sys.exit(1) @@ -164,8 +164,8 @@ def run_playbook_file( debug: bool = False ): """Requests server to execute a playbook from local path.""" - url = f"{base_url}/playbooks/execute/file" - logger.info(f"Requesting server execute playbook file: {playbook_file_path_on_server}") + url = f'{base_url}/playbooks/execute/file' + logger.info(f'Requesting server execute playbook file: {playbook_file_path_on_server}') payload = {'file_path': playbook_file_path_on_server} try: params = {'debug': True} if debug else {} @@ -183,13 +183,13 @@ def run_playbook_file( if not data.get('success'): sys.exit(1) except httpx.RequestError as e: - logger.error(f"HTTP Request Error executing playbook file: {e}") + logger.error(f'HTTP Request Error executing playbook file: {e}') sys.exit(1) except httpx.HTTPStatusError as e: - logger.error(f"HTTP Status Error (File): {e.response.status_code} - {e.response.text}") + logger.error(f'HTTP Status Error (File): {e.response.status_code} - {e.response.text}') sys.exit(1) except Exception as e: - logger.error(f"Unexpected error (File): {e}", exc_info=True) + logger.error(f'Unexpected error (File): {e}', exc_info=True) sys.exit(1) @@ -219,14 +219,14 @@ def run_command(client: httpx.Client, base_url: str, args): body_dict[pydantic_field_name] = arg_value try: - logger.debug(f"Sending POST to {url}") - logger.debug(f"Request Body: {json.dumps(body_dict, indent=2)}") + logger.debug(f'Sending POST to {url}') + logger.debug(f'Request Body: {json.dumps(body_dict, indent=2)}') response = client.post(url, json=body_dict, headers=get_auth_headers()) response.raise_for_status() data = response.json() update_token_from_response(data) - logger.info(f"Received response from /{type} endpoint.") - logger.debug(f"Response data: {data}") + logger.info(f'Received response from /{type} endpoint.') + logger.debug(f'Response data: {data}') result = data.get('result', {}) state = data.get('state', {}).get('variables', {}) @@ -246,13 +246,13 @@ def run_command(client: httpx.Client, base_url: str, args): sys.exit(1) except httpx.RequestError as e: - logger.error(f"HTTP Request Error executing command: {e}") + logger.error(f'HTTP Request Error executing command: {e}') sys.exit(1) except httpx.HTTPStatusError as e: - logger.error(f"HTTP Status Error ({url}): {e.response.status_code} - {e.response.text}") + logger.error(f'HTTP Status Error ({url}): {e.response.status_code} - {e.response.text}') sys.exit(1) except Exception as e: - logger.error(f"Unexpected error executing command: {e}", exc_info=True) + logger.error(f'Unexpected error executing command: {e}', exc_info=True) sys.exit(1) @@ -362,9 +362,9 @@ def main(): if args.cacert: cert_path = os.path.abspath(args.cacert) # Ensure absolute path if os.path.exists(cert_path): - logger.info(f"Configured httpx to verify using CA cert: {cert_path}") + logger.info(f'Configured httpx to verify using CA cert: {cert_path}') else: - logger.error(f"CA certificate file not found at specified path: {cert_path}") + logger.error(f'CA certificate file not found at specified path: {cert_path}') sys.exit(1) # Create HTTP Client @@ -388,12 +388,12 @@ def main(): sys.exit(1) except httpx.ConnectError as e: logger.error( - f"Connection Error: Could not connect to {args.base_url}. " - f"Is the server running with HTTPS? Did you provide cert? Details: {e}" + f'Connection Error: Could not connect to {args.base_url}. ' + f'Is the server running with HTTPS? Did you provide cert? Details: {e}' ) sys.exit(1) except Exception as main_err: - logger.error(f"Client execution failed: {main_err}", exc_info=True) + logger.error(f'Client execution failed: {main_err}', exc_info=True) sys.exit(1) logger.info('Client finished.') diff --git a/remote_rest/create_hashes.py b/remote_rest/create_hashes.py index c073f928..4aa479e5 100644 --- a/remote_rest/create_hashes.py +++ b/remote_rest/create_hashes.py @@ -1,4 +1,3 @@ -import os from passlib.context import CryptContext diff --git a/remote_rest/main.py b/remote_rest/main.py index 8af51041..cd20ed0b 100644 --- a/remote_rest/main.py +++ b/remote_rest/main.py @@ -52,8 +52,8 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: # any other async startup tasks ? except Exception as e: - logger.critical(f"Failed to initialize during startup lifespan: {e}", exc_info=True) - raise RuntimeError(f"Failed to initialize application state: {e}") from e + logger.critical(f'Failed to initialize during startup lifespan: {e}', exc_info=True) + raise RuntimeError(f'Failed to initialize application state: {e}') from e yield # Application runs here @@ -63,13 +63,13 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: for instance_id in instance_ids: instance = state.INSTANCES.pop(instance_id, None) if instance: - logger.info(f"Cleaning up instance {instance_id}...") + logger.info(f'Cleaning up instance {instance_id}...') try: # blocking? instance.clean_session_stores() instance.pm.kill_or_wait_processes() except Exception as e: - logger.error(f"Error cleaning up instance {instance_id}: {e}", exc_info=True) + logger.error(f'Error cleaning up instance {instance_id}: {e}', exc_info=True) logger.info('Instance cleanup complete (lifespan).') @@ -83,7 +83,7 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: # Exception Handling @app.exception_handler(ExecException) async def attackmate_execution_exception_handler(request: Request, exc: ExecException): - logger.error(f"AttackMate Execution Exception: {exc}") + logger.error(f'AttackMate Execution Exception: {exc}') return JSONResponse( status_code=400, content={ @@ -97,14 +97,14 @@ async def attackmate_execution_exception_handler(request: Request, exc: ExecExce @app.exception_handler(Exception) async def generic_exception_handler(request: Request, exc: Exception): if isinstance(exc, SystemExit): - logger.error(f"Command triggered SystemExit with code {exc.code}") + logger.error(f'Command triggered SystemExit with code {exc.code}') return JSONResponse( status_code=400, # client-side error pattern content={ 'detail': 'Command execution led to termination request', 'error_message': ( f"SystemExit triggered (likely due to error condition like 'exit_on_error'). " - f"Exit code: {exc.code}" + f'Exit code: {exc.code}' ), 'instance_id': None }, @@ -117,7 +117,7 @@ async def generic_exception_handler(request: Request, exc: Exception): @app.post('/login', response_model=TokenResponse, tags=['Auth']) async def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends()): """Authenticates user and returns an access token.""" - logger.info(f"Login attempt for user: {form_data.username}") + logger.info(f'Login attempt for user: {form_data.username}') hashed_password = get_user_hash(form_data.username) if not hashed_password: logger.warning(f"Login failed: User '{form_data.username}' not found.") @@ -154,10 +154,10 @@ async def root(): if __name__ == '__main__': if not os.path.exists(KEY_FILE): - logger.critical(f"SSL Error: Key file not found at {KEY_FILE}") + logger.critical(f'SSL Error: Key file not found at {KEY_FILE}') sys.exit(1) if not os.path.exists(CERT_FILE): - logger.critical(f"SSL Error: Certificate file not found at {CERT_FILE}") + logger.critical(f'SSL Error: Certificate file not found at {CERT_FILE}') sys.exit(1) uvicorn.run('remote_rest.main:app', host='0.0.0.0', diff --git a/remote_rest/routers/commands.py b/remote_rest/routers/commands.py index ec703f01..3a5f5cef 100644 --- a/remote_rest/routers/commands.py +++ b/remote_rest/routers/commands.py @@ -5,9 +5,9 @@ from attackmate.schemas.base import BaseCommand from fastapi import APIRouter, Depends, Header, HTTPException -from src.attackmate.execexception import ExecException -from src.attackmate.result import Result as AttackMateResult -from src.attackmate.schemas.command_types import Command +from attackmate.execexception import ExecException +from attackmate.result import Result as AttackMateResult +from attackmate.schemas.command_types import Command from remote_rest.auth_utils import API_KEY_HEADER_NAME, get_current_user @@ -31,14 +31,14 @@ async def run_command_on_instance(instance: AttackMate, command_data: BaseComman logger.info(f"Executing command type '{command_data.type}' on instance") # type: ignore # TODO does this work? need to pass command class object here? result = instance.run_command(command_data) - logger.info(f"Command execution finished. RC: {result.returncode}") + logger.info(f'Command execution finished. RC: {result.returncode}') return result except (ExecException, SystemExit) as e: - logger.error(f"AttackMate execution error: {e}", exc_info=True) + logger.error(f'AttackMate execution error: {e}', exc_info=True) raise e except Exception as e: - logger.error(f"Unexpected error during instance.run_command: {e}", exc_info=True) - raise HTTPException(status_code=500, detail=f"Internal server error during command execution: {e}") + logger.error(f'Unexpected error during instance.run_command: {e}', exc_info=True) + raise HTTPException(status_code=500, detail=f'Internal server error during command execution: {e}') @router.post('/execute', response_model=ExecutionResponseModel) diff --git a/remote_rest/routers/playbooks.py b/remote_rest/routers/playbooks.py index 8e75d14e..dcc1f486 100644 --- a/remote_rest/routers/playbooks.py +++ b/remote_rest/routers/playbooks.py @@ -31,7 +31,7 @@ def read_log_file(log_path: Optional[str]) -> Optional[str]: return f.read() except Exception as e: logger.error(f"Failed to read log file '{log_path}': {e}") - return f"Error reading log file: {e}" + return f'Error reading log file: {e}' # Playbook Execution @@ -42,8 +42,8 @@ async def execute_playbook_from_yaml(playbook_yaml: str = Body(..., media_type=' False, description="Enable debug logging for this request's instance log." ), - current_user: str = Depends(get_current_user), - x_auth_token: Optional[str] = Header(None, alias=API_KEY_HEADER_NAME)): + current_user: str = Depends(get_current_user), + x_auth_token: Optional[str] = Header(None, alias=API_KEY_HEADER_NAME)): """ Executes a playbook provided as YAML content in the request body. Use a transient AttackMate instance. @@ -58,20 +58,20 @@ async def execute_playbook_from_yaml(playbook_yaml: str = Body(..., media_type=' if not playbook_dict: raise ValueError('Received empty or invalid playbook YAML content.') playbook = Playbook.model_validate(playbook_dict) - logger.info(f"Creating transient AttackMate instance, ID: {instance_id}") + logger.info(f'Creating transient AttackMate instance, ID: {instance_id}') am_instance = AttackMate(playbook=playbook, config=attackmate_config, varstore=None) return_code = am_instance.main() final_state = varstore_to_state_model(am_instance.varstore) - logger.info(f"Transient playbook execution finished. return code {return_code}") + logger.info(f'Transient playbook execution finished. return code {return_code}') attackmate_log = read_log_file(attackmate_log_path) output_log = read_log_file(output_log_path) json_log = read_log_file(json_log_path) except (yaml.YAMLError, ValidationError, ValueError) as e: - logger.error(f"Playbook validation/parsing error: {e}") - raise HTTPException(status_code=422, detail=f"Invalid playbook YAML: {e}") + logger.error(f'Playbook validation/parsing error: {e}') + raise HTTPException(status_code=422, detail=f'Invalid playbook YAML: {e}') except Exception as e: - logger.error(f"Unexpected error during playbook execution: {e}", exc_info=True) - raise HTTPException(status_code=500, detail=f"Server error during playbook execution: {e}") + logger.error(f'Unexpected error during playbook execution: {e}', exc_info=True) + raise HTTPException(status_code=500, detail=f'Server error during playbook execution: {e}') finally: if am_instance: logger.info('Cleaning up transient playbook instance.') @@ -79,7 +79,7 @@ async def execute_playbook_from_yaml(playbook_yaml: str = Body(..., media_type=' am_instance.clean_session_stores() am_instance.pm.kill_or_wait_processes() except Exception as cleanup_e: - logger.error(f"Error cleaning transient instance: {cleanup_e}", exc_info=True) + logger.error(f'Error cleaning transient instance: {cleanup_e}', exc_info=True) return PlaybookResponseModel( success=(return_code == 0), @@ -100,7 +100,7 @@ async def execute_playbook_from_file(request_body: PlaybookFileRequest, description=( "Enable debug level logging for this request's instance log." ) - ), + ), current_user: str = Depends(get_current_user), x_auth_token: Optional[str] = Header(None, alias=API_KEY_HEADER_NAME) ): @@ -109,7 +109,7 @@ async def execute_playbook_from_file(request_body: PlaybookFileRequest, Uses a transient AttackMate instance. """ # TODO ensure this only executes playbooks in certain locations -> read up on path traversal - logger.info(f"Received request to execute playbook from file: {request_body.file_path}") + logger.info(f'Received request to execute playbook from file: {request_body.file_path}') try: # base directory exists if not os.path.isdir(ALLOWED_PLAYBOOK_DIR): @@ -130,13 +130,13 @@ async def execute_playbook_from_file(request_body: PlaybookFileRequest, # Check if the file exists if not os.path.isfile(full_path): - raise FileNotFoundError(f"Playbook file not found atpath: {full_path}") + raise FileNotFoundError(f'Playbook file not found atpath: {full_path}') except (ValueError, FileNotFoundError) as e: - logger.error(f"Invalid or non-existent playbook path requested: {request_body.file_path} -> {e}") - raise HTTPException(status_code=400, detail=f"Invalid or non-existent playbook file path: {e}") + logger.error(f'Invalid or non-existent playbook path requested: {request_body.file_path} -> {e}') + raise HTTPException(status_code=400, detail=f'Invalid or non-existent playbook file path: {e}') except Exception as e: - logger.error(f"Error processing playbook path: {e}", exc_info=True) + logger.error(f'Error processing playbook path: {e}', exc_info=True) raise HTTPException(status_code=500, detail='Server error processing file path.') instance_id = str(uuid.uuid4()) @@ -144,13 +144,13 @@ async def execute_playbook_from_file(request_body: PlaybookFileRequest, with instance_logging(instance_id, log_level) as log_files: attackmate_log_path, output_log_path, json_log_path = log_files try: - logger.info(f"Parsing playbook from: {full_path}") + logger.info(f'Parsing playbook from: {full_path}') playbook = parse_playbook(full_path, logger) - logger.info(f"Creating transient AttackMate instance, ID: {instance_id}") + logger.info(f'Creating transient AttackMate instance, ID: {instance_id}') am_instance = AttackMate(playbook=playbook, config=attackmate_config, varstore=None) return_code = am_instance.main() final_state = varstore_to_state_model(am_instance.varstore) - logger.info(f"Transient playbook execution finished. RC: {return_code}") + logger.info(f'Transient playbook execution finished. RC: {return_code}') attackmate_log = read_log_file(attackmate_log_path) output_log = read_log_file(output_log_path) json_log = read_log_file(json_log_path) @@ -159,8 +159,8 @@ async def execute_playbook_from_file(request_body: PlaybookFileRequest, raise HTTPException( status_code=400, detail=f"Invalid playbook content in file '{request_body.file_path}': {e}") except Exception as e: - logger.error(f"Unexpected error during playbook file execution: {e}", exc_info=True) - raise HTTPException(status_code=500, detail=f"Server error during playbook execution: {e}") + logger.error(f'Unexpected error during playbook file execution: {e}', exc_info=True) + raise HTTPException(status_code=500, detail=f'Server error during playbook execution: {e}') finally: if am_instance: logger.info('Cleaning up transient playbook instance.') @@ -168,7 +168,7 @@ async def execute_playbook_from_file(request_body: PlaybookFileRequest, am_instance.clean_session_stores() am_instance.pm.kill_or_wait_processes() except Exception as e: - logger.error(f"Cleanup error: {e}") + logger.error(f'Cleanup error: {e}') return PlaybookResponseModel( success=(return_code == 0), diff --git a/src/attackmate/executors/browser/browserexecutor.py b/src/attackmate/executors/browser/browserexecutor.py index 33d1d57e..8d77aea6 100644 --- a/src/attackmate/executors/browser/browserexecutor.py +++ b/src/attackmate/executors/browser/browserexecutor.py @@ -13,7 +13,7 @@ def __init__(self, pm, varstore, **kwargs): def log_command(self, command: BrowserCommand): self.logger.info( - f"Executing Browser Command: {command.cmd}" + f'Executing Browser Command: {command.cmd}' f"{f' {command.selector}' if command.selector else ''}" f"{f' {command.url}' if command.url else ''}" ) @@ -35,7 +35,7 @@ def _exec_cmd(self, command: BrowserCommand) -> Result: if self.session_store.has_session(command.creates_session): self.logger.warning( f"Session '{command.creates_session}' already exists! " - f"Closing it before creating a new one." + f'Closing it before creating a new one.' ) self.session_store.close_session(command.creates_session) session_thread = SessionThread( @@ -57,7 +57,7 @@ def _exec_cmd(self, command: BrowserCommand) -> Result: elif command.cmd == 'screenshot': session_thread.submit_command('screenshot', screenshot_path=command.screenshot_path) else: - return Result(f"Unknown browser command: {command.cmd}", 1) + return Result(f'Unknown browser command: {command.cmd}', 1) return Result('Browser command executed successfully.', 0) diff --git a/src/attackmate/executors/browser/sessionstore.py b/src/attackmate/executors/browser/sessionstore.py index f2056936..98945de6 100644 --- a/src/attackmate/executors/browser/sessionstore.py +++ b/src/attackmate/executors/browser/sessionstore.py @@ -8,6 +8,7 @@ class SessionThread(threading.Thread): A dedicated thread that manages a single Playwright browser session (browser, context, page). All commands for this session are queued to make sure they run synchronously in this thread. """ + def __init__(self, session_name=None, headless=False): super().__init__() self.session_name = session_name @@ -100,7 +101,7 @@ def _handle_command(self, cmd_name, *args, **kwargs): screenshot_path = kwargs['screenshot_path'] self.page.screenshot(path=screenshot_path) else: - raise ValueError(f"Unknown command: {cmd_name}") + raise ValueError(f'Unknown command: {cmd_name}') return 'OK' @@ -129,6 +130,7 @@ class BrowserSessionStore: """ Manages named sessions. Each named session has its own SessionThread. """ + def __init__(self): self.sessions = {} self.lock = threading.Lock() diff --git a/src/attackmate/executors/common/debugexecutor.py b/src/attackmate/executors/common/debugexecutor.py index 4246598d..477c9bfd 100644 --- a/src/attackmate/executors/common/debugexecutor.py +++ b/src/attackmate/executors/common/debugexecutor.py @@ -22,7 +22,7 @@ def _exec_cmd(self, command: DebugCommand) -> Result: self.setoutputvars = False ret = 0 if command.wait_for_key: - self.logger.warning("Type enter to continue") + self.logger.warning('Type enter to continue') input() if command.exit: ret = 1 diff --git a/src/attackmate/executors/features/cmdvars.py b/src/attackmate/executors/features/cmdvars.py index 0e50a11f..853de6be 100644 --- a/src/attackmate/executors/features/cmdvars.py +++ b/src/attackmate/executors/features/cmdvars.py @@ -36,11 +36,11 @@ def substitute_template_vars(self, command: BaseCommand, substitute_cmd_vars: bo for attr_name in command.list_template_vars(): # Skip variable substitution for "cmd" - if attr_name == 'cmd' and not substitute_cmd_vars: - continue + if attr_name == 'cmd' and not substitute_cmd_vars: + continue - if attr_name == "break_if": - continue + if attr_name == 'break_if': + continue attr_value = getattr(command, attr_name) diff --git a/src/attackmate/executors/features/exitonerror.py b/src/attackmate/executors/features/exitonerror.py index 8d66393c..052620d0 100644 --- a/src/attackmate/executors/features/exitonerror.py +++ b/src/attackmate/executors/features/exitonerror.py @@ -23,8 +23,8 @@ def error_if(self, command: BaseCommand, result: Result): m = re.search(command.error_if, result.stdout, re.MULTILINE) if m is not None: self.logger.error( - f'Exitting because error_if matches: {m.group(0)}' - ) + f'Exitting because error_if matches: {m.group(0)}' + ) exit(1) def error_if_not(self, command: BaseCommand, result: Result): @@ -32,6 +32,6 @@ def error_if_not(self, command: BaseCommand, result: Result): m = re.search(command.error_if_not, result.stdout, re.MULTILINE) if m is None: self.logger.error( - 'Exitting because error_if_not does not match' - ) + 'Exitting because error_if_not does not match' + ) exit(1) diff --git a/src/attackmate/executors/metasploit/msfexecutor.py b/src/attackmate/executors/metasploit/msfexecutor.py index 89dc516a..5fd8787b 100644 --- a/src/attackmate/executors/metasploit/msfexecutor.py +++ b/src/attackmate/executors/metasploit/msfexecutor.py @@ -10,6 +10,7 @@ from attackmate.schemas.metasploit import MsfModuleCommand from attackmate.executors.metasploit.msfsessionstore import MsfSessionStore from attackmate.processmanager import ProcessManager +from multiprocessing.managers import BaseManager from multiprocessing import Manager from multiprocessing.queues import JoinableQueue from attackmate.executors.executor_factory import executor_factory @@ -29,6 +30,7 @@ def __init__( self.msfconfig = msfconfig self.sessionstore = msfsessionstore self.msf = None + self.manager: Optional[BaseManager] = None super().__init__(pm, varstore, cmdconfig) def _create_queue(self) -> Optional[JoinableQueue]: @@ -140,12 +142,12 @@ def cleanup(self): if active_sessions: for session_id, session_data in active_sessions.items(): try: - self.logger.debug(f"Stopping msf session {session_id}") + self.logger.debug(f'Stopping msf session {session_id}') self.msf.sessions.session(session_id).stop() - self.logger.info(f"Msf session {session_id} stopped successfully.") + self.logger.info(f'Msf session {session_id} stopped successfully.') except Exception as e: - self.logger.error(f"Failed to stop msf session {session_id}: {str(e)}") + self.logger.error(f'Failed to stop msf session {session_id}: {str(e)}') else: - self.logger.debug("No active msf sessions found.") + self.logger.debug('No active msf sessions found.') -#msf.session returs shell manager +# msf.session returs shell manager diff --git a/src/attackmate/executors/metasploit/msfsessionexecutor.py b/src/attackmate/executors/metasploit/msfsessionexecutor.py index 1b860ea4..b32bfd6f 100644 --- a/src/attackmate/executors/metasploit/msfsessionexecutor.py +++ b/src/attackmate/executors/metasploit/msfsessionexecutor.py @@ -92,10 +92,10 @@ def cleanup(self): if active_sessions: for session_id, session_data in active_sessions.items(): try: - self.logger.debug(f"Stopping msf session {session_id}") + self.logger.debug(f'Stopping msf session {session_id}') self.msf.sessions.session(session_id).stop() - self.logger.info(f"Msf session {session_id} stopped successfully.") + self.logger.info(f'Msf session {session_id} stopped successfully.') except Exception as e: - self.logger.error(f"Failed to stop msf session {session_id}: {str(e)}") + self.logger.error(f'Failed to stop msf session {session_id}: {str(e)}') else: self.logger.debug('No active msf sessions found.') diff --git a/src/attackmate/executors/metasploit/msfsessionstore.py b/src/attackmate/executors/metasploit/msfsessionstore.py index 9c4fa3b3..bf2bdc22 100644 --- a/src/attackmate/executors/metasploit/msfsessionstore.py +++ b/src/attackmate/executors/metasploit/msfsessionstore.py @@ -34,9 +34,10 @@ def get_session_by_name(self, name: str, msfsessions, block: bool = True) -> str if v['exploit_uuid'] == self.sessions[name]: return k else: - self.logger.debug(f"uuid {self.sessions[name]} does not match with any entry in sessions") + self.logger.debug( + f'uuid {self.sessions[name]} does not match with any entry in sessions') else: - self.logger.debug(f"{name} not found in msfsessions") + self.logger.debug(f'{name} not found in msfsessions') if not block: raise ExecException(f'Session ({name}) not found') diff --git a/src/attackmate/executors/sliver/sliversessionexecutor.py b/src/attackmate/executors/sliver/sliversessionexecutor.py index dadeee19..7e522f9e 100644 --- a/src/attackmate/executors/sliver/sliversessionexecutor.py +++ b/src/attackmate/executors/sliver/sliversessionexecutor.py @@ -291,7 +291,7 @@ async def cleanup(self): await self.client.kill_beacon(beacon.ID) except Exception as e: self.logger.error(f'Error cleaning up sliver sessions: {e}') - + def _exec_cmd(self, command: SliverSessionCommand) -> Result: loop = asyncio.get_event_loop() diff --git a/src/attackmate/executors/ssh/sessionstore.py b/src/attackmate/executors/ssh/sessionstore.py index a253767a..83f46367 100644 --- a/src/attackmate/executors/ssh/sessionstore.py +++ b/src/attackmate/executors/ssh/sessionstore.py @@ -50,7 +50,6 @@ def set_existing_session(self, session_name: str, if self.has_session(session_name): self.set_session(session_name, client, channel) - def clean_sessions(self): """ Closes all active SSH sessions and their associated channels in the session store, @@ -74,4 +73,3 @@ def clean_sessions(self): self.logger.error(f"Error closing client for ssh session '{session_name}': {e}") self.store.clear() - diff --git a/src/attackmate/executors/ssh/sshexecutor.py b/src/attackmate/executors/ssh/sshexecutor.py index 13b6b791..5f1b5631 100644 --- a/src/attackmate/executors/ssh/sshexecutor.py +++ b/src/attackmate/executors/ssh/sshexecutor.py @@ -121,7 +121,7 @@ def connect_use_session(self, command: SFTPCommand | SSHCommand) -> SSHClient: if command.creates_session is not None: self.session_store.set_session(command.creates_session, client) return client - + def cleanup(self): self.session_store.clean_sessions() diff --git a/src/attackmate/executors/vnc/sessionstore.py b/src/attackmate/executors/vnc/sessionstore.py index d2b95111..88dc164c 100644 --- a/src/attackmate/executors/vnc/sessionstore.py +++ b/src/attackmate/executors/vnc/sessionstore.py @@ -35,4 +35,3 @@ def clean_sessions(self): self.logger.error(f"Error closing vnc client for session '{session_name}': {e}") self.store.clear() - diff --git a/src/attackmate/executors/vnc/vncexecutor.py b/src/attackmate/executors/vnc/vncexecutor.py index 1d292643..42ffcd07 100644 --- a/src/attackmate/executors/vnc/vncexecutor.py +++ b/src/attackmate/executors/vnc/vncexecutor.py @@ -46,9 +46,9 @@ def build_connection_string( raise ExecException('Hostname is required for VNC connection.') connection_str = hostname if display is not None: - connection_str += f":{display}" + connection_str += f':{display}' elif port is not None: - connection_str += f"::{port}" + connection_str += f'::{port}' return connection_str def connect(self, command: VncCommand) -> api.ThreadedVNCClientProxy: @@ -64,10 +64,10 @@ def connect(self, command: VncCommand) -> api.ThreadedVNCClientProxy: start_time = time.time() while time.time() - start_time < self.connection_timeout: if client and client.protocol and client.protocol.connected: - self.logger.info(f"Connected to VNC server: {connection_string}") + self.logger.info(f'Connected to VNC server: {connection_string}') return client time.sleep(0.1) # Poll every 100ms for connection status - self.logger.info(f"Could not connect to VNC server: {connection_string}") + self.logger.info(f'Could not connect to VNC server: {connection_string}') client.disconnect() return None @@ -85,7 +85,7 @@ def connect_use_session(self, command): # If 'session' is specified, check if it exists, else raise an error if not self.session_store.has_session(command.session): raise ExecException( - f"VNC-Session not in Session-Store: {command.session}" + f'VNC-Session not in Session-Store: {command.session}' ) else: return self.session_store.get_client_by_session(command.session) @@ -135,16 +135,16 @@ def _exec_cmd(self, command: VncCommand) -> Result: action() else: - raise ExecException(f"Unknown VNC command: {command.cmd}") + raise ExecException(f'Unknown VNC command: {command.cmd}') except TimeoutError: self.cleanup() raise ExecException( - f"VNC Timeout Error after {command.expect_timeout} seconds" + f'VNC Timeout Error after {command.expect_timeout} seconds' ) except (ValueError, AttributeError, AuthenticationError, OSError) as e: - raise ExecException(f"VNC Execution Error: {e}") + raise ExecException(f'VNC Execution Error: {e}') output = 'vnc_connected' return Result(output, 0) @@ -156,7 +156,7 @@ def close_connection(self, session_name: str = 'default'): if client: try: self.logger.info( - f"Closing VNC connection for session: {session_name}" + f'Closing VNC connection for session: {session_name}' ) client.disconnect() api.shutdown() @@ -168,7 +168,7 @@ def close_connection(self, session_name: str = 'default'): self.logger.error( f"Error closing VNC session '{session_name}': {str(e)}" ) - raise ExecException(f"Error closing VNC connection: {e}") + raise ExecException(f'Error closing VNC connection: {e}') else: self.logger.warning( f"VNC session '{session_name}' not found in session store." diff --git a/src/attackmate/result.py b/src/attackmate/result.py index 3fde98ba..bb0ad697 100644 --- a/src/attackmate/result.py +++ b/src/attackmate/result.py @@ -26,4 +26,4 @@ def __init__(self, stdout, returncode): self.returncode = returncode def __repr__(self): - return f"Result(stdout={repr(self.stdout)}, returncode={self.returncode})" + return f'Result(stdout={repr(self.stdout)}, returncode={self.returncode})' diff --git a/src/attackmate/schemas/command_types.py b/src/attackmate/schemas/command_types.py index 4a54728f..8bc53113 100644 --- a/src/attackmate/schemas/command_types.py +++ b/src/attackmate/schemas/command_types.py @@ -1,18 +1,12 @@ - - -from typing import List, Annotated, TypeAlias, Union from pydantic import Field +from typing import List, TypeAlias, Union, Annotated from attackmate.schemas.command_subtypes import RemotelyExecutableCommand from attackmate.schemas.remote import AttackMateRemoteCommand -Command: TypeAlias = Annotated[ - Union[ - RemotelyExecutableCommand, - AttackMateRemoteCommand - ], - Field(discriminator='type'), -] - +Command: TypeAlias = Annotated[Union[ + RemotelyExecutableCommand, + AttackMateRemoteCommand +], Field(discriminator='type')] -Commands: TypeAlias = List[Command] +Commands: TypeAlias = List[Command] diff --git a/test/units/test_browserexecutor.py b/test/units/test_browserexecutor.py index a642a519..b3f2bb4a 100644 --- a/test/units/test_browserexecutor.py +++ b/test/units/test_browserexecutor.py @@ -7,7 +7,7 @@ # Minimal, stable inline HTML (no network needed) -HTML_SIMPLE = "

Hello World

" +HTML_SIMPLE = '

Hello World

' HTML_WITH_LINK = """ @@ -18,8 +18,8 @@ """ -DATA_URL_SIMPLE = "data:text/html," + quote(HTML_SIMPLE) -DATA_URL_WITH_LINK = "data:text/html," + quote(HTML_WITH_LINK) +DATA_URL_SIMPLE = 'data:text/html,' + quote(HTML_SIMPLE) +DATA_URL_WITH_LINK = 'data:text/html,' + quote(HTML_WITH_LINK) @pytest.fixture diff --git a/test/units/test_commanddelay.py b/test/units/test_commanddelay.py index 228ecdc7..26896add 100644 --- a/test/units/test_commanddelay.py +++ b/test/units/test_commanddelay.py @@ -34,11 +34,11 @@ def test_command_delay_is_applied(): expected_maximum_time = expected_minimum_time + 0.5 assert elapsed_time >= expected_minimum_time, ( - f"Execution faster ({elapsed_time:.4f}s) than the minimum expected delay " - f"({expected_minimum_time:.4f}s)." + f'Execution faster ({elapsed_time:.4f}s) than the minimum expected delay ' + f'({expected_minimum_time:.4f}s).' ) assert elapsed_time < expected_maximum_time, ( - f"Execution slower ({elapsed_time:.4f}s) than expected." + f'Execution slower ({elapsed_time:.4f}s) than expected.' ) @@ -62,7 +62,7 @@ def test_zero_command_delay(): # With no delay, execution should be very fast. assert elapsed_time < 0.1, ( - f"Execution with no delay took too long: {elapsed_time:.4f}s." + f'Execution with no delay took too long: {elapsed_time:.4f}s.' ) @@ -87,5 +87,5 @@ def test_delay_is_not_applied_for_exempt_commands(): elapsed_time = end_time - start_time assert elapsed_time < 0.1, ( - f"Execution with exempt commands took too long: {elapsed_time:.4f}s." + f'Execution with exempt commands took too long: {elapsed_time:.4f}s.' ) diff --git a/test/units/test_loopexecutor.py b/test/units/test_loopexecutor.py index 2dd96ddc..460d59d9 100644 --- a/test/units/test_loopexecutor.py +++ b/test/units/test_loopexecutor.py @@ -48,7 +48,6 @@ def caplog_setup(caplog): return caplog - class TestLoopExecutor: def test_items(self, caplog_setup, varstore, loop_executor): @@ -74,14 +73,18 @@ def test_range(self, caplog_setup, varstore, loop_executor): def test_break_if_with_range(self, caplog_setup, loop_executor): caplog = caplog_setup lc = LoopCommand( - type='loop', cmd='range(1,5)', break_if='$LOOP_INDEX =~ 2', commands=[DebugCommand(cmd='$LOOP_INDEX', type='debug')] - ) + type='loop', + cmd='range(1,5)', + break_if='$LOOP_INDEX =~ 2', + commands=[ + DebugCommand( + cmd='$LOOP_INDEX', + type='debug')]) loop_executor.run(lc) assert 'Debug: \'1\'' in [rec.message for rec in caplog.records] assert 'Debug: \'2\'' not in [rec.message for rec in caplog.records] assert 'Debug: \'3\'' not in [rec.message for rec in caplog.records] - def test_until(self, caplog_setup, loop_executor): caplog = caplog_setup lc = LoopCommand( diff --git a/test/units/test_vncexecutor.py b/test/units/test_vncexecutor.py index 80a5b8c7..72c0aa3a 100644 --- a/test/units/test_vncexecutor.py +++ b/test/units/test_vncexecutor.py @@ -1,11 +1,9 @@ import pytest -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock from attackmate.executors.vnc.vncexecutor import VncExecutor from attackmate.schemas.vnc import VncCommand from attackmate.variablestore import VariableStore from attackmate.processmanager import ProcessManager -from attackmate.execexception import ExecException -from vncdotool.client import AuthenticationError @pytest.fixture @@ -41,7 +39,7 @@ def test_vnc_connect_success(vnc_executor, mock_vnc_client, mocker): mock_vnc_client.protocol.connected = True # Mock successful connection result = vnc_executor._exec_cmd(command) - + assert result.stdout == 'vnc_connected' mock_vnc_client.keyPress.assert_called_once_with('a') @@ -75,5 +73,3 @@ def test_vnc_create_and_use_session(vnc_executor, mock_vnc_client, mocker): assert result.stdout == 'vnc_connected' mock_vnc_client.keyPress.assert_any_call('a') mock_vnc_client.keyPress.assert_any_call('b') - -