From b2cc5fac415a1211ee0dc2ee6f808cc62eeccd3f Mon Sep 17 00:00:00 2001
From: Ilija Lazoroski <ilija.la@live.com>
Date: Fri, 10 Mar 2023 17:22:25 +0100
Subject: [PATCH] SMB: Add SMBOptions to SMBPlugin

Issue: #2952
PR: #3089
---
 .../exploiters/smb/src/smb_options.py         | 27 ++++++++
 .../exploiters/smb/test_smb_options.py        | 62 +++++++++++++++++++
 2 files changed, 89 insertions(+)
 create mode 100644 monkey/agent_plugins/exploiters/smb/src/smb_options.py
 create mode 100644 monkey/tests/unit_tests/agent_plugins/exploiters/smb/test_smb_options.py

diff --git a/monkey/agent_plugins/exploiters/smb/src/smb_options.py b/monkey/agent_plugins/exploiters/smb/src/smb_options.py
new file mode 100644
index 00000000000..3d9b45832ba
--- /dev/null
+++ b/monkey/agent_plugins/exploiters/smb/src/smb_options.py
@@ -0,0 +1,27 @@
+from pydantic import Field
+
+from common.base_models import InfectionMonkeyBaseModel
+
+
+class SMBOptions(InfectionMonkeyBaseModel):
+    agent_binary_upload_timeout: float = Field(
+        default=30.0,
+        gt=0.0,
+        le=100.0,
+        description="Maximum time allowed for uploading the Agent binary to the target.",
+    )
+    use_kerberos: bool = Field(
+        default=False, description="Should the RPC transport use Kerberos authentication."
+    )
+    rpc_connect_timeout: float = Field(
+        default=15.0,
+        gt=0.0,
+        le=100.0,
+        description="The maximum time (in seconds) to wait for RPC connection.",
+    )
+    smb_connect_timeout: float = Field(
+        default=15.0,
+        gt=0.0,
+        le=100.0,
+        description="The maximum time (in seconds) to wait for SMB connection.",
+    )
diff --git a/monkey/tests/unit_tests/agent_plugins/exploiters/smb/test_smb_options.py b/monkey/tests/unit_tests/agent_plugins/exploiters/smb/test_smb_options.py
new file mode 100644
index 00000000000..a01c5968f9e
--- /dev/null
+++ b/monkey/tests/unit_tests/agent_plugins/exploiters/smb/test_smb_options.py
@@ -0,0 +1,62 @@
+import pydantic
+import pytest
+from agent_plugins.exploiters.smb.src.smb_options import SMBOptions
+
+UPLOAD_TIMEOUT = 12.4
+USE_KERBEROS = True
+RPC_CONNECT_TIMEOUT = 11.2
+SMB_CONNECT_TIMEOUT = 42.1
+
+
+SMB_OPTIONS_DICT = {
+    "agent_binary_upload_timeout": UPLOAD_TIMEOUT,
+    "use_kerberos": USE_KERBEROS,
+    "rpc_connect_timeout": RPC_CONNECT_TIMEOUT,
+    "smb_connect_timeout": SMB_CONNECT_TIMEOUT,
+}
+
+
+SMB_OPTIONS_OBJECT = SMBOptions(
+    agent_binary_upload_timeout=UPLOAD_TIMEOUT,
+    use_kerberos=USE_KERBEROS,
+    rpc_connect_timeout=RPC_CONNECT_TIMEOUT,
+    smb_connect_timeout=SMB_CONNECT_TIMEOUT,
+)
+
+UPLOAD_TIMEOUT_EXCEPTION = {"agent_binary_upload_timeout": 70000}
+RPC_TIMEOUT_EXCEPTION = {"rpc_connect_timeout": -1}
+SMB_TIMEOUT_EXCEPTION = {"smb_connect_timeout": 101}
+
+
+def test_smb_options__serialization():
+    assert SMB_OPTIONS_OBJECT.dict(simplify=True) == SMB_OPTIONS_DICT
+
+
+def test_smb_options__full_serialization():
+    assert SMBOptions(**SMB_OPTIONS_OBJECT.dict(simplify=True)) == SMB_OPTIONS_OBJECT
+
+
+def test_smb_options__deserialization():
+    assert SMBOptions(**SMB_OPTIONS_DICT) == SMB_OPTIONS_OBJECT
+
+
+def test_smb_options__default():
+    smb_options = SMBOptions()
+
+    assert smb_options.agent_binary_upload_timeout == 30.0
+    assert smb_options.use_kerberos is False
+    assert smb_options.rpc_connect_timeout == 15.0
+    assert smb_options.smb_connect_timeout == 15.0
+
+
+@pytest.mark.parametrize(
+    "options_dict",
+    [
+        UPLOAD_TIMEOUT_EXCEPTION,
+        RPC_TIMEOUT_EXCEPTION,
+        SMB_TIMEOUT_EXCEPTION,
+    ],
+)
+def test_smb_options_constrains(options_dict):
+    with pytest.raises((pydantic.errors.NumberNotLeError, pydantic.errors.NumberNotGtError)):
+        SMBOptions(**options_dict)