From f118253634158b00a6011da4aadfac9293b02f9b Mon Sep 17 00:00:00 2001 From: thorinaboenke Date: Tue, 8 Sep 2026 21:05:37 +0200 Subject: [PATCH 1/3] add extra forbidden error and suggestions --- src/attackmate/playbook_parser.py | 25 ++++++++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/src/attackmate/playbook_parser.py b/src/attackmate/playbook_parser.py index a045a616..16b1380e 100644 --- a/src/attackmate/playbook_parser.py +++ b/src/attackmate/playbook_parser.py @@ -1,14 +1,25 @@ import os +import difflib import logging import traceback import yaml -from typing import Optional +from typing import Any, List, Optional from pathlib import Path from pydantic import ValidationError +from attackmate.command import CommandRegistry from attackmate.schemas.playbook import Playbook from attackmate.schemas.config import Config +def _valid_fields_for_error(error: Any, playbook_yaml: Any) -> Optional[List[str]]: + try: + command = playbook_yaml['commands'][error['loc'][1]] + model = CommandRegistry.get_command_class(command['type'], command.get('cmd')) + except (KeyError, IndexError, TypeError, ValueError): + return None + return list(model.model_fields.keys()) + + def load_configfile(config_file: str) -> Config: with open(config_file) as f: config = yaml.safe_load(f) @@ -153,5 +164,17 @@ def parse_playbook(playbook_file: str, logger: logging.Logger) -> Playbook: f'Value error in command {int(error["loc"][-2]) + 1}: ' f'{error["loc"][-1]} - {error["msg"]}' ) + elif error['type'] == 'extra_forbidden': + bad_key = str(error['loc'][-1]) + cmd_type = error['loc'][-2] + valid_fields = _valid_fields_for_error(error, playbook_yaml) + suggestion = '' + if valid_fields: + matches = difflib.get_close_matches(bad_key, valid_fields, n=1) + if matches: + suggestion = f" , did you mean '{matches[0]}'?" + logger.error( + f"Unknown field in {cmd_type} command: '{bad_key}'{suggestion}" + ) logger.error(traceback.format_exc()) exit(1) From e51bddcbeb0c87fe02231ec34dc50d7f92acb574 Mon Sep 17 00:00:00 2001 From: thorinaboenke Date: Tue, 8 Sep 2026 21:07:37 +0200 Subject: [PATCH 2/3] correct registry type hints --- src/attackmate/command.py | 8 ++++---- src/attackmate/schemas/base.py | 5 ++++- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/src/attackmate/command.py b/src/attackmate/command.py index d851f0e4..bbe096af 100644 --- a/src/attackmate/command.py +++ b/src/attackmate/command.py @@ -1,4 +1,4 @@ -from typing import Dict, Tuple, Optional +from typing import Dict, Tuple, Optional, Type from attackmate.schemas.base import BaseCommand @@ -13,8 +13,8 @@ class CommandRegistry: The more specific ``(type, cmd)`` key takes precedence over the ``type``-only key. """ - _type_registry: Dict[str, BaseCommand] = {} - _type_cmd_registry: Dict[Tuple[str, str], BaseCommand] = {} + _type_registry: Dict[str, Type[BaseCommand]] = {} + _type_cmd_registry: Dict[Tuple[str, str], Type[BaseCommand]] = {} @classmethod def register(cls, type_: str, cmd: Optional[str] = None): @@ -32,7 +32,7 @@ class ShellCommand(BaseCommand): ... """ - def decorator(command_class: BaseCommand): + def decorator(command_class: Type[BaseCommand]): if cmd: cls._type_cmd_registry[(type_, cmd)] = command_class else: diff --git a/src/attackmate/schemas/base.py b/src/attackmate/schemas/base.py index 70b4e4c0..7d6253f1 100644 --- a/src/attackmate/schemas/base.py +++ b/src/attackmate/schemas/base.py @@ -1,6 +1,6 @@ from typing import Annotated, List, Optional, Dict from dataclasses import field -from pydantic import AfterValidator, BeforeValidator, BaseModel, ValidationInfo +from pydantic import AfterValidator, BeforeValidator, BaseModel, ConfigDict, ValidationInfo import re # https://stackoverflow.com/questions/71539448/using-different-pydantic-models-depending-on-the-value-of-fields @@ -27,6 +27,9 @@ def check_var_pattern(value: str, info: ValidationInfo) -> str: class BaseCommand(BaseModel): + # Reject unknown/misspelled command parameters + model_config = ConfigDict(extra='forbid') + def list_template_vars(self) -> List[str]: """Get a list of all variables that can be used as templates From eb4408784d0c40e588018d6e2c14d68220a547e8 Mon Sep 17 00:00:00 2001 From: thorinaboenke Date: Tue, 8 Sep 2026 21:13:46 +0200 Subject: [PATCH 3/3] add test for misspelled parameters --- test/units/test_parseplaybook.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/test/units/test_parseplaybook.py b/test/units/test_parseplaybook.py index c288d19b..12bec6b9 100644 --- a/test/units/test_parseplaybook.py +++ b/test/units/test_parseplaybook.py @@ -56,3 +56,26 @@ def test_parse_playbook_nonexistent_file(mock_logger): with pytest.raises(SystemExit): parse_playbook(str(non_existent_file), mock_logger) mock_logger.error.assert_called_with(error_message) + + +misspelled_param_playbook_yaml = """ +### +commands: + - type: shell + cmd: ls + backgound: true +""" + + +def test_parse_playbook_misspelled_param(mock_logger, tmp_path): + playbook_file = tmp_path / 'playbook.yml' + playbook_file.write_text(misspelled_param_playbook_yaml) + + with pytest.raises(SystemExit): + parse_playbook(str(playbook_file), mock_logger) + + messages = [call.args[0] for call in mock_logger.error.call_args_list] + assert any( + "Unknown field in shell command: 'backgound'" in m and "did you mean 'background'?" in m + for m in messages + ), messages