Skip to content

Commit

Permalink
SMB: Add SMBOptions to SMBPlugin
Browse files Browse the repository at this point in the history
Issue: #2952
PR: #3089
  • Loading branch information
ilija-lazoroski committed Mar 13, 2023
1 parent 818959d commit b2cc5fa
Show file tree
Hide file tree
Showing 2 changed files with 89 additions and 0 deletions.
27 changes: 27 additions & 0 deletions monkey/agent_plugins/exploiters/smb/src/smb_options.py
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.",
)
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)

0 comments on commit b2cc5fa

Please sign in to comment.