From 2342a9f60d6ec0d851572c201efc35d2d7882738 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Tue, 16 May 2023 15:36:18 +0000 Subject: [PATCH 1/3] Common: Add sha256 to AgentRegistrationData --- monkey/common/agent_registration_data.py | 1 + .../island_api_client/test_http_island_api_client.py | 2 ++ .../island_event_handlers/test_handle_agent_registration.py | 5 ++++- .../cc/models/test_agent_registration_message.py | 5 +++++ .../unit_tests/monkey_island/cc/resources/test_agents.py | 1 + 5 files changed, 13 insertions(+), 1 deletion(-) diff --git a/monkey/common/agent_registration_data.py b/monkey/common/agent_registration_data.py index 5d281be25bc..2be163bcb0c 100644 --- a/monkey/common/agent_registration_data.py +++ b/monkey/common/agent_registration_data.py @@ -17,6 +17,7 @@ class AgentRegistrationData(InfectionMonkeyBaseModel): parent_id: Optional[UUID] cc_server: SocketAddress network_interfaces: Sequence[IPv4Interface] + sha256: str _make_immutable_sequence = validator("network_interfaces", pre=True, allow_reuse=True)( make_immutable_sequence diff --git a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py index dac957cc71e..3200cca595e 100644 --- a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py +++ b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py @@ -42,6 +42,7 @@ SERVER = SocketAddress(ip="1.1.1.1", port=9999) WINDOWS = "windows" AGENT_ID = UUID("80988359-a1cd-42a2-9b47-5b94b37cd673") +AGENT_SHA256 = "7f295ebba73b2f549f98b9c35ae588f153283476ac8f087cc75d9c4788795d82" AGENT_REGISTRATION = AgentRegistrationData( id=AGENT_ID, machine_hardware_id=1, @@ -49,6 +50,7 @@ parent_id=None, cc_server=SERVER, network_interfaces=[], + sha256=AGENT_SHA256, ) TIMESTAMP = 123456789 diff --git a/monkey/tests/unit_tests/monkey_island/cc/island_event_handlers/test_handle_agent_registration.py b/monkey/tests/unit_tests/monkey_island/cc/island_event_handlers/test_handle_agent_registration.py index f0bd28b0b84..e6a10f60658 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/island_event_handlers/test_handle_agent_registration.py +++ b/monkey/tests/unit_tests/monkey_island/cc/island_event_handlers/test_handle_agent_registration.py @@ -19,6 +19,7 @@ ) AGENT_ID = UUID("860aff5b-d2af-43ea-afb5-62bac3d30b7e") +AGENT_SHA256 = "a4beff584bc18ef48f64874e2f57ef2c8088d6947d4e15490120730401640dbc" SEED_ID = 10 @@ -38,6 +39,7 @@ parent_id=None, cc_server=SocketAddress.from_string(IP), network_interfaces=[IPv4Interface("192.168.1.2/24")], + sha256=AGENT_SHA256, ) @@ -67,7 +69,6 @@ def node_repository() -> INodeRepository: @pytest.fixture def handler(machine_repository, agent_repository, node_repository) -> handle_agent_registration: - return handle_agent_registration( machine_repository, agent_repository, node_repository, get_current_datetime=lambda: NOW ) @@ -125,6 +126,7 @@ def test_existing_machine_updated__find_by_ip(handler, machine_repository): IPv4Interface("192.168.1.4/24"), IPv4Interface("192.168.1.5/24"), ], + sha256=AGENT_SHA256, ) existing_machine = Machine( @@ -230,6 +232,7 @@ def test_machine_interfaces_updated(handler, machine_repository): IPv4Interface("192.168.1.3/16"), IPv4Interface("192.168.1.4/24"), ], + sha256=AGENT_SHA256, ) expected_network_interfaces = sorted( (*agent_registration_data.network_interfaces, existing_machine.network_interfaces[-1]) diff --git a/monkey/tests/unit_tests/monkey_island/cc/models/test_agent_registration_message.py b/monkey/tests/unit_tests/monkey_island/cc/models/test_agent_registration_message.py index dff49c252e6..a09ac686841 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/models/test_agent_registration_message.py +++ b/monkey/tests/unit_tests/monkey_island/cc/models/test_agent_registration_message.py @@ -10,6 +10,7 @@ AGENT_ID = UUID("012e7238-7b81-4108-8c7f-0787bc3f3c10") PARENT_ID = UUID("0fc9afcb-1902-436b-bd5c-1ad194252484") +AGENT_SHA256 = "c6fef92294c27bbe372696bd58f35d1117e41088edff3ffd3c73df865297ac88" SOCKET_ADDRESS = SocketAddress(ip="192.168.1.1", port=5000) AGENT_REGISTRATION_MESSAGE_OBJECT_DICT = { @@ -19,6 +20,7 @@ "parent_id": PARENT_ID, "cc_server": SOCKET_ADDRESS, "network_interfaces": [IPv4Interface("10.0.0.1/24"), IPv4Interface("192.168.5.32/16")], + "sha256": AGENT_SHA256, } AGENT_REGISTRATION_MESSAGE_SIMPLE_DICT = { @@ -28,6 +30,7 @@ "parent_id": str(PARENT_ID), "cc_server": SOCKET_ADDRESS.dict(simplify=True), "network_interfaces": ["10.0.0.1/24", "192.168.5.32/16"], + "sha256": AGENT_SHA256, } @@ -54,6 +57,7 @@ def test_from_serialized(): ("parent_id", 2.1), ("cc_server", [1]), ("network_interfaces", "not-a-list"), + ("sha256", []), ], ) def test_construct_invalid_field__type_error(key, value): @@ -90,6 +94,7 @@ def test_construct_invalid_field__value_error(key, value): ("parent_id", AGENT_ID), ("cc_server", SOCKET_ADDRESS), ("network_interfaces", ["10.0.0.1/24"]), + ("sha256", "fdf59a0b2fc2b49b9e245da81a44654e97fa5be566cc719deeb77aaf74fe6013"), ], ) def test_fields_immutable(key, value): diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agents.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agents.py index a49bd82324d..a7096be592e 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agents.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agents.py @@ -19,6 +19,7 @@ "parent_id": UUID("9d55ba33-95c2-417d-bd86-d3d11e47daeb"), "cc_server": {"ip": "10.0.0.1", "port": "5000"}, "network_interfaces": ["10.1.1.2/24"], + "sha256": "cf5c10a8073aa923877ee66df8c1912cac2dbb4b85a97d09cb95d57bde4d2876", } AGENTS = ( From fd03a5f452844b605b2faa3d693584ad7b20c039 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Tue, 16 May 2023 15:37:17 +0000 Subject: [PATCH 2/3] Agent: Register a dummy value for sha256 --- monkey/infection_monkey/monkey.py | 1 + 1 file changed, 1 insertion(+) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index e7d75734913..2cfd12d6c44 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -248,6 +248,7 @@ def _register_agent(self): parent_id=self._opts.parent, cc_server=self._island_address, network_interfaces=get_network_interfaces(), + sha256="0000000000000000000000000000000000000000000000000000000000000000", ) self._island_api_client.register_agent(agent_registration_data) From b3a95074d4f53591c634fc1e8766f4e0a79fda11 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Tue, 16 May 2023 16:40:37 +0000 Subject: [PATCH 3/3] Common: Add validation to AgentRegistrationData.sha256 field --- monkey/common/agent_registration_data.py | 4 ++-- .../cc/models/test_agent_registration_message.py | 3 +++ 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/monkey/common/agent_registration_data.py b/monkey/common/agent_registration_data.py index 2be163bcb0c..239fea731d8 100644 --- a/monkey/common/agent_registration_data.py +++ b/monkey/common/agent_registration_data.py @@ -3,7 +3,7 @@ from typing import Optional, Sequence from uuid import UUID -from pydantic import validator +from pydantic import Field, validator from .base_models import InfectionMonkeyBaseModel from .transforms import make_immutable_sequence @@ -17,7 +17,7 @@ class AgentRegistrationData(InfectionMonkeyBaseModel): parent_id: Optional[UUID] cc_server: SocketAddress network_interfaces: Sequence[IPv4Interface] - sha256: str + sha256: str = Field(regex=r"^[0-9a-fA-F]{64}$") _make_immutable_sequence = validator("network_interfaces", pre=True, allow_reuse=True)( make_immutable_sequence diff --git a/monkey/tests/unit_tests/monkey_island/cc/models/test_agent_registration_message.py b/monkey/tests/unit_tests/monkey_island/cc/models/test_agent_registration_message.py index a09ac686841..5352b8fc6d5 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/models/test_agent_registration_message.py +++ b/monkey/tests/unit_tests/monkey_island/cc/models/test_agent_registration_message.py @@ -75,6 +75,9 @@ def test_construct_invalid_field__type_error(key, value): ("start_time", "not-a-date-time"), ("network_interfaces", [1, "stuff", 3]), ("cc_server", []), + ("sha256", "not-a-hex-string-although-it-is-of-the-correct-length-for-sha256"), + ("sha256", "12345678123456781234567812345678123456781234567812345678123456780"), + ("sha256", "12345678"), ], ) def test_construct_invalid_field__value_error(key, value):