From 554a180fbd18305beb417b1812e70ce1e124f083 Mon Sep 17 00:00:00 2001
From: Shreya Malviya <shreya.malviya@gmail.com>
Date: Mon, 29 Aug 2022 13:24:13 +0530
Subject: [PATCH] Common: Create CustomPBAConfiguration using pydantic

---
 .../agent_sub_configurations.py               | 38 +++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/monkey/common/agent_configuration/agent_sub_configurations.py b/monkey/common/agent_configuration/agent_sub_configurations.py
index 4ed94d7a8..be1c9ca87 100644
--- a/monkey/common/agent_configuration/agent_sub_configurations.py
+++ b/monkey/common/agent_configuration/agent_sub_configurations.py
@@ -1,6 +1,12 @@
 from dataclasses import dataclass
 from typing import Dict, Tuple
 
+from pydantic import validator
+
+from common.base_models import MutableInfectionMonkeyBaseModel
+
+from .validators import validate_linux_filename, validate_windows_filename
+
 
 @dataclass(frozen=True)
 class CustomPBAConfiguration:
@@ -25,6 +31,38 @@ class CustomPBAConfiguration:
     windows_filename: str
 
 
+class Pydantic___CustomPBAConfiguration(MutableInfectionMonkeyBaseModel):
+    """
+    A configuration for custom post-breach actions
+
+    Attributes:
+        :param linux_command: Command to run on Linux victim machines. If a file is uploaded,
+                              use this field to change its permissions, execute it, and/or delete it
+                              Example: `chmod +x file.sh; ./file.sh; rm file.sh`
+        :param linux_filename: Name of the file to upload on Linux victim machines
+        :param windows_command: Command to run on Windows victim machines. If a file is uploaded,
+                                use this field to change its permissions, execute it, and/or delete
+                                it
+                                Example: `file.bat & del file.bat`
+        :param windows_filename: Name of the file to upload on Windows victim machines
+    """
+
+    linux_command: str
+    linux_filename: str
+    windows_command: str
+    windows_filename: str
+
+    @validator("linux_filename")
+    def linux_filename_valid(cls, filename):
+        validate_linux_filename(filename)
+        return filename
+
+    @validator("windows_filename")
+    def windows_filename_valid(cls, filename):
+        validate_windows_filename(filename)
+        return filename
+
+
 @dataclass(frozen=True)
 class PluginConfiguration:
     """