diff --git a/monkey/monkey_island/cc/models/__init__.py b/monkey/monkey_island/cc/models/__init__.py index 31f2700d0..94e063a81 100644 --- a/monkey/monkey_island/cc/models/__init__.py +++ b/monkey/monkey_island/cc/models/__init__.py @@ -12,3 +12,4 @@ from .user_credentials import UserCredentials from .machine import Machine, MachineID from .communication_type import CommunicationType from .node import Node +from .agent import Agent diff --git a/monkey/monkey_island/cc/models/agent.py b/monkey/monkey_island/cc/models/agent.py new file mode 100644 index 000000000..73ff8d693 --- /dev/null +++ b/monkey/monkey_island/cc/models/agent.py @@ -0,0 +1,18 @@ +from datetime import datetime +from typing import Optional +from uuid import UUID + +from pydantic import Field + +from . import MachineID +from .base_models import MutableBaseModel + + +class Agent(MutableBaseModel): + id: UUID = Field(..., allow_mutation=False) + machine_id: MachineID = Field(..., allow_mutation=False) + start_time: datetime = Field(..., allow_mutation=False) + stop_time: Optional[datetime] + parent_id: UUID = Field(..., allow_mutation=False) + cc_server: str = Field(default="") + log_contents: str = Field(default="") diff --git a/monkey/tests/unit_tests/monkey_island/cc/models/test_agent.py b/monkey/tests/unit_tests/monkey_island/cc/models/test_agent.py new file mode 100644 index 000000000..1851a4873 --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/models/test_agent.py @@ -0,0 +1,133 @@ +from datetime import datetime, timezone +from uuid import UUID + +import pytest + +from monkey_island.cc.models import Agent + +AGENT_ID = UUID("012e7238-7b81-4108-8c7f-0787bc3f3c10") +PARENT_ID = UUID("0fc9afcb-1902-436b-bd5c-1ad194252484") + +AGENT_OBJECT_DICT = { + "id": AGENT_ID, + "machine_id": 2, + "parent_id": PARENT_ID, + "start_time": datetime.fromtimestamp(1660848408, tz=timezone.utc), +} + +AGENT_SIMPLE_DICT = { + "id": str(AGENT_ID), + "machine_id": 2, + "parent_id": str(PARENT_ID), + "start_time": "2022-08-18T18:46:48+00:00", +} + + +def test_constructor__defaults_from_objects(): + a = Agent(**AGENT_OBJECT_DICT) + + assert a.stop_time is None + assert a.cc_server == "" + + +def test_constructor__defaults_from_simple_dict(): + a = Agent(**AGENT_SIMPLE_DICT) + + assert a.stop_time is None + assert a.cc_server == "" + assert a.log_contents == "" + + +def test_to_dict(): + a = Agent(**AGENT_OBJECT_DICT) + agent_simple_dict = AGENT_SIMPLE_DICT.copy() + agent_simple_dict["stop_time"] = None + agent_simple_dict["cc_server"] = "" + agent_simple_dict["log_contents"] = "" + + assert a.dict(simplify=True) == agent_simple_dict + + +@pytest.mark.parametrize( + "key, value", + [ + ("id", 1), + ("machine_id", "not-an-int"), + ("start_time", None), + ("stop_time", []), + ("parent_id", 2.1), + ("cc_server", []), + ("log_contents", None), + ], +) +def test_construct_invalid_field__type_error(key, value): + invalid_type_dict = AGENT_SIMPLE_DICT.copy() + invalid_type_dict[key] = value + + with pytest.raises(TypeError): + Agent(**invalid_type_dict) + + +@pytest.mark.parametrize( + "key, value", + [ + ("machine_id", -1), + ("start_time", "not-a-datetime"), + ("stop_time", "not-a-datetime"), + ], +) +def test_construct_invalid_field__value_error(key, value): + invalid_value_dict = AGENT_SIMPLE_DICT.copy() + invalid_value_dict[key] = value + + with pytest.raises(ValueError): + Agent(**invalid_value_dict) + + +def test_id_immutable(): + a = Agent(**AGENT_SIMPLE_DICT) + + with pytest.raises(TypeError): + a.id = PARENT_ID + + +def test_machine_id_immutable(): + a = Agent(**AGENT_SIMPLE_DICT) + + with pytest.raises(TypeError): + a.machine_id = 10 + + +def test_start_time_immutable(): + a = Agent(**AGENT_SIMPLE_DICT) + + with pytest.raises(TypeError): + a.start_time = 100 + + +def test_parent_id_immutable(): + a = Agent(**AGENT_SIMPLE_DICT) + + with pytest.raises(TypeError): + a.parent_id = AGENT_ID + + +def test_stop_time_set_validated(): + a = Agent(**AGENT_SIMPLE_DICT) + + with pytest.raises(ValueError): + a.stop_time = "testing!" + + +def test_cc_server_set_validated(): + a = Agent(**AGENT_SIMPLE_DICT) + + with pytest.raises(ValueError): + a.cc_server = None + + +def test_log_contents_set_validated(): + a = Agent(**AGENT_SIMPLE_DICT) + + with pytest.raises(ValueError): + a.log_contents = None