Merge pull request #79 from ait-testbed/fix_var_int

Automatically Convert int-variables to type str
ait-testbed · Jun 14, 2024 · bc3e725 · bc3e725
2 parents eb9ad00 + 126e3e9
commit bc3e725
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/src/attackmate/schemas/base.py b/src/attackmate/schemas/base.py
@@ -21,6 +21,8 @@ def check_var_pattern(value: str, info: ValidationInfo) -> str:
                          BeforeValidator(transform_int_to_str),
                          AfterValidator(check_var_pattern)]
 
+# Like StringNumber but without checks for valid $variable
+StrInt = Annotated[Optional[str | int], BeforeValidator(transform_int_to_str)]
 
 class BaseCommand(BaseModel):
     def list_template_vars(self) -> List[str]:

diff --git a/src/attackmate/schemas/playbook.py b/src/attackmate/schemas/playbook.py
@@ -27,6 +27,7 @@
 from .tempfile import TempfileCommand
 from .debug import DebugCommand
 from .regex import RegExCommand
+from .base import StrInt
 
 
 Commands = List[Union[
@@ -62,5 +63,5 @@
 
 
 class Playbook(BaseModel):
-    vars: Optional[Dict[str, str]] = None
+    vars: Optional[Dict[str, StrInt]] = None
     commands: Commands