-
Notifications
You must be signed in to change notification settings - Fork 786
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
818959d
commit b2cc5fa
Showing
2 changed files
with
89 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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.", | ||
) |
62 changes: 62 additions & 0 deletions
62
monkey/tests/unit_tests/agent_plugins/exploiters/smb/test_smb_options.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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) |