From 126e3e94b9bfbf8cfa73a508e659ba73c1f66707 Mon Sep 17 00:00:00 2001
From: Wolfgang Hotwagner <wolfgang.hotwagner@ait.ac.at>
Date: Fri, 14 Jun 2024 10:46:56 +0200
Subject: [PATCH] Automatically Convert int-variables to type str

---
 src/attackmate/schemas/base.py     | 2 ++
 src/attackmate/schemas/playbook.py | 3 ++-
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/attackmate/schemas/base.py b/src/attackmate/schemas/base.py
index 6d52fef..cea4d0b 100644
--- 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
index 701298a..cb2bf35 100644
--- 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