Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 21 additions & 2 deletions .pre-commit-config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
29 changes: 11 additions & 18 deletions docs/source/developing/command.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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('<command_type>')`` decorator.

For example, to add a ``debug`` command:
Expand All @@ -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']
Expand All @@ -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.


Expand All @@ -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('<command_type>')`` decorator.
Accordingly, executors must be registered using the ``@executor_factory.register_executor('<command_type>')`` 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.

::
Expand Down Expand Up @@ -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
=====================================================
Expand All @@ -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.







2 changes: 1 addition & 1 deletion docs/source/developing/integration.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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}")
print(f"Command failed with return code {result.returncode}")
2 changes: 1 addition & 1 deletion docs/source/installation/uv.rst
Original file line number Diff line number Diff line change
Expand Up @@ -59,4 +59,4 @@ Installation Steps
.. warning::

Please note that you need to :ref:`sliver-fix` if you want
to use the sliver commands!
to use the sliver commands!
5 changes: 0 additions & 5 deletions docs/source/playbook/commands/vnc.rst
Original file line number Diff line number Diff line change
Expand Up @@ -151,8 +151,3 @@ Execute commands on a remote server via VNC. Uses the `vncdotool <https://github
.. note::

The vnc connection needs to be closed with the command ``close`` explicitely, otherwise attackmate will keep running.





3 changes: 0 additions & 3 deletions examples/regex.yml
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,3 @@ commands:

- type: debug
cmd: "Result string: $SUBSTITUTED"



3 changes: 3 additions & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -52,3 +52,6 @@ where = ["src"]

[tool.setuptools.dynamic]
version = {attr = "attackmate.metadata.__version__"}

[tool.mypy]
explicit_package_bases = true
6 changes: 3 additions & 3 deletions remote_rest/auth_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ def verify_password(plain_password: str, hashed_password: str) -> 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)


Expand Down Expand Up @@ -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']
Expand All @@ -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
70 changes: 35 additions & 35 deletions remote_rest/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down Expand Up @@ -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
Expand All @@ -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()
Expand All @@ -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)


Expand All @@ -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 {}
Expand All @@ -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)


Expand Down Expand Up @@ -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', {})
Expand All @@ -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)


Expand Down Expand Up @@ -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
Expand All @@ -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.')
Expand Down
1 change: 0 additions & 1 deletion remote_rest/create_hashes.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
import os

from passlib.context import CryptContext

Expand Down
Loading