Merge branch 'prototype-register-agents' into develop

CHANGELOG.md

@@ -23,6 +23,7 @@ Changelog](https://keepachangelog.com/en/1.0.0/).
 - `/api/events` endpoint. #2155
 - The ability to customize the file extension used by ransomware when
   encrypting files. #1242
+- `/api/agents` endpoint.
 
 ### Changed
 - Reset workflow. Now it's possible to delete data gathered by agents without

monkey/common/__init__.py

@@ -1,5 +1,9 @@
 """
 Used for a common things between agent and island
 """
+from . import transforms
 from .di_container import DIContainer, UnresolvableDependencyError
 from .operating_system import OperatingSystem
+from . import types
+from . import base_models
+from .agent_registration_data import AgentRegistrationData

monkey/common/agent_registration_data.py

@@ -0,0 +1,23 @@
+from datetime import datetime
+from ipaddress import IPv4Interface
+from typing import Optional, Sequence
+from uuid import UUID
+
+from pydantic import validator
+
+from .base_models import InfectionMonkeyBaseModel
+from .transforms import make_immutable_sequence
+from .types import HardwareID
+
+
+class AgentRegistrationData(InfectionMonkeyBaseModel):
+    id: UUID
+    machine_hardware_id: HardwareID
+    start_time: datetime
+    parent_id: Optional[UUID]
+    cc_server: str
+    network_interfaces: Sequence[IPv4Interface]
+
+    _make_immutable_sequence = validator("network_interfaces", pre=True, allow_reuse=True)(
+        make_immutable_sequence
+    )

monkey/common/base_models.py

@@ -5,6 +5,7 @@ from pydantic import BaseModel, Extra, ValidationError
 
 
 class InfectionMonkeyModelConfig:
+    allow_mutation = False
     underscore_attrs_are_private = True
     extra = Extra.forbid
 

monkey/common/types.py

@@ -0,0 +1,3 @@
+from pydantic import PositiveInt
+
+HardwareID = PositiveInt

monkey/infection_monkey/Pipfile

@@ -27,6 +27,7 @@ paramiko = {editable = true, ref = "2.10.3.dev1", git = "https://github.com/Vaka
 marshmallow = "*"
 marshmallow-enum = "*"
 pypubsub = "*"
+pydantic = "*"
 
 [dev-packages]
 ldap3 = "*"

monkey/infection_monkey/Pipfile.lock

@@ -1,7 +1,7 @@
 {
     "_meta": {
         "hash": {
-            "sha256": "4b34d3b744fa5d28b36d6f1dae271474b6a92db3a221b8c14d3c04eedfab5d9d"
+            "sha256": "0ae0a7c88cba4dbd3ad91fd472f6bf12399a9819931d1bf3a936623fa2bfcb6d"
         },
         "pipfile-spec": 6,
         "requires": {
@@ -630,6 +630,47 @@
             "markers": "python_version >= '2.7' and python_version not in '3.0, 3.1, 3.2, 3.3, 3.4'",
             "version": "==3.15.0"
         },
+        "pydantic": {
+            "hashes": [
+                "sha256:1061c6ee6204f4f5a27133126854948e3b3d51fcc16ead2e5d04378c199b2f44",
+                "sha256:19b5686387ea0d1ea52ecc4cffb71abb21702c5e5b2ac626fd4dbaa0834aa49d",
+                "sha256:2bd446bdb7755c3a94e56d7bdfd3ee92396070efa8ef3a34fab9579fe6aa1d84",
+                "sha256:328558c9f2eed77bd8fffad3cef39dbbe3edc7044517f4625a769d45d4cf7555",
+                "sha256:32e0b4fb13ad4db4058a7c3c80e2569adbd810c25e6ca3bbd8b2a9cc2cc871d7",
+                "sha256:3ee0d69b2a5b341fc7927e92cae7ddcfd95e624dfc4870b32a85568bd65e6131",
+                "sha256:4aafd4e55e8ad5bd1b19572ea2df546ccace7945853832bb99422a79c70ce9b8",
+                "sha256:4b3946f87e5cef3ba2e7bd3a4eb5a20385fe36521d6cc1ebf3c08a6697c6cfb3",
+                "sha256:4de71c718c9756d679420c69f216776c2e977459f77e8f679a4a961dc7304a56",
+                "sha256:5565a49effe38d51882cb7bac18bda013cdb34d80ac336428e8908f0b72499b0",
+                "sha256:5803ad846cdd1ed0d97eb00292b870c29c1f03732a010e66908ff48a762f20e4",
+                "sha256:5da164119602212a3fe7e3bc08911a89db4710ae51444b4224c2382fd09ad453",
+                "sha256:615661bfc37e82ac677543704437ff737418e4ea04bef9cf11c6d27346606044",
+                "sha256:78a4d6bdfd116a559aeec9a4cfe77dda62acc6233f8b56a716edad2651023e5e",
+                "sha256:7d0f183b305629765910eaad707800d2f47c6ac5bcfb8c6397abdc30b69eeb15",
+                "sha256:7ead3cd020d526f75b4188e0a8d71c0dbbe1b4b6b5dc0ea775a93aca16256aeb",
+                "sha256:84d76ecc908d917f4684b354a39fd885d69dd0491be175f3465fe4b59811c001",
+                "sha256:8cb0bc509bfb71305d7a59d00163d5f9fc4530f0881ea32c74ff4f74c85f3d3d",
+                "sha256:91089b2e281713f3893cd01d8e576771cd5bfdfbff5d0ed95969f47ef6d676c3",
+                "sha256:9c9e04a6cdb7a363d7cb3ccf0efea51e0abb48e180c0d31dca8d247967d85c6e",
+                "sha256:a8c5360a0297a713b4123608a7909e6869e1b56d0e96eb0d792c27585d40757f",
+                "sha256:afacf6d2a41ed91fc631bade88b1d319c51ab5418870802cedb590b709c5ae3c",
+                "sha256:b34ba24f3e2d0b39b43f0ca62008f7ba962cff51efa56e64ee25c4af6eed987b",
+                "sha256:bd67cb2c2d9602ad159389c29e4ca964b86fa2f35c2faef54c3eb28b4efd36c8",
+                "sha256:c0f5e142ef8217019e3eef6ae1b6b55f09a7a15972958d44fbd228214cede567",
+                "sha256:cdb4272678db803ddf94caa4f94f8672e9a46bae4a44f167095e4d06fec12979",
+                "sha256:d70916235d478404a3fa8c997b003b5f33aeac4686ac1baa767234a0f8ac2326",
+                "sha256:d8ce3fb0841763a89322ea0432f1f59a2d3feae07a63ea2c958b2315e1ae8adb",
+                "sha256:e0b214e57623a535936005797567231a12d0da0c29711eb3514bc2b3cd008d0f",
+                "sha256:e631c70c9280e3129f071635b81207cad85e6c08e253539467e4ead0e5b219aa",
+                "sha256:e78578f0c7481c850d1c969aca9a65405887003484d24f6110458fb02cca7747",
+                "sha256:f0ca86b525264daa5f6b192f216a0d1e860b7383e3da1c65a1908f9c02f42801",
+                "sha256:f1a68f4f65a9ee64b6ccccb5bf7e17db07caebd2730109cb8a95863cfa9c4e55",
+                "sha256:fafe841be1103f340a24977f61dee76172e4ae5f647ab9e7fd1e1fca51524f08",
+                "sha256:ff68fc85355532ea77559ede81f35fff79a6a5543477e168ab3a381887caea76"
+            ],
+            "index": "pypi",
+            "version": "==1.9.2"
+        },
         "pyinstaller": {
             "hashes": [
                 "sha256:24035eb9fffa2e3e288b4c1c9710043819efc7203cae5c8c573bec16f4a8e98f",
@@ -935,7 +976,6 @@
                 "sha256:1d6b085e5c445141c475476000b661f60fff1aaa19f76bf82b7abb92e0ff4942",
                 "sha256:b6a6be5711b1b6c8d55bda7a8befd75c48c12b770b9d227d31c1737dbf0d40a6"
             ],
-            "index": "pypi",
             "markers": "sys_platform == 'win32'",
             "version": "==1.5.1"
         },

monkey/infection_monkey/i_control_channel.py

@@ -1,11 +1,23 @@
 import abc
-from typing import Sequence
+from typing import Optional, Sequence
+from uuid import UUID
 
 from common.agent_configuration import AgentConfiguration
 from common.credentials import Credentials
 
 
 class IControlChannel(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def register_agent(self, parent_id: Optional[UUID] = None):
+        """
+        Registers this agent with the Island when this agent starts
+
+        :param parent: The ID of the parent that spawned this agent, or None if this agent has no
+                       parent
+        :raises IslandCommunicationError: If the agent cannot be successfully registered
+        """
+        pass
+
     @abc.abstractmethod
     def should_agent_stop(self) -> bool:
         """

monkey/infection_monkey/master/control_channel.py

@@ -1,14 +1,18 @@
 import json
 import logging
 from pprint import pformat
-from typing import Mapping, Sequence
+from typing import Mapping, Optional, Sequence
+from uuid import UUID
 
 import requests
 
+from common import AgentRegistrationData
 from common.agent_configuration import AgentConfiguration
 from common.common_consts.timeouts import SHORT_REQUEST_TIMEOUT
 from common.credentials import Credentials
 from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError
+from infection_monkey.utils import agent_process
+from infection_monkey.utils.ids import get_agent_id, get_machine_id
 
 requests.packages.urllib3.disable_warnings()
 
@@ -21,6 +25,34 @@ class ControlChannel(IControlChannel):
         self._control_channel_server = server
         self._proxies = proxies
 
+    def register_agent(self, parent: Optional[UUID] = None):
+        agent_registration_data = AgentRegistrationData(
+            id=get_agent_id(),
+            machine_hardware_id=get_machine_id(),
+            start_time=agent_process.get_start_time(),
+            parent_id=parent,
+            cc_server=self._control_channel_server,
+            network_interfaces=[],  # TODO: Populate this
+        )
+
+        try:
+            url = f"https://{self._control_channel_server}/api/agents"
+            response = requests.post(  # noqa: DUO123
+                url,
+                json=agent_registration_data.dict(simplify=True),
+                verify=False,
+                proxies=self._proxies,
+                timeout=SHORT_REQUEST_TIMEOUT,
+            )
+            response.raise_for_status()
+        except (
+            requests.exceptions.ConnectionError,
+            requests.exceptions.Timeout,
+            requests.exceptions.TooManyRedirects,
+            requests.exceptions.HTTPError,
+        ) as e:
+            raise IslandCommunicationError(e)
+
     def should_agent_stop(self) -> bool:
         if not self._control_channel_server:
             logger.error("Agent should stop because it can't connect to the C&C server.")

monkey/infection_monkey/monkey.py

@@ -175,6 +175,7 @@ class InfectionMonkey:
         control_channel = ControlChannel(
             self._control_client.server_address, GUID, self._control_client.proxies
         )
+        control_channel.register_agent(self._opts.parent)
 
         config = control_channel.get_config()
         self._monkey_inbound_tunnel = self._control_client.create_control_tunnel(

monkey/infection_monkey/utils/ids.py

@@ -0,0 +1,22 @@
+from uuid import UUID, getnode, uuid4
+
+
+def get_agent_id() -> UUID:
+    """
+    Get the agent ID for the current running agent
+
+    Each time an agent process starts, the return value of this function will be unique. Subsequent
+    calls to this function from within the same process will have the same return value.
+    """
+    if get_agent_id._id is None:
+        get_agent_id._id = uuid4()
+
+    return get_agent_id._id
+
+
+get_agent_id._id = None
+
+
+def get_machine_id() -> int:
+    """Get an integer that uniquely defines the machine the agent is running on"""
+    return getnode()

monkey/monkey_island/cc/app.py

@@ -13,6 +13,7 @@ from monkey_island.cc.database import database, mongo
 from monkey_island.cc.resources import (
     AgentBinaries,
     AgentConfiguration,
+    Agents,
     ClearSimulationData,
     Events,
     IPAddresses,
@@ -162,6 +163,7 @@ def init_restful_endpoints(api: FlaskDIWrapper):
     api.add_resource(Register)
     api.add_resource(RegistrationStatus)
     api.add_resource(Authenticate)
+    api.add_resource(Agents)
     api.add_resource(Monkey)
     api.add_resource(LocalRun)
     api.add_resource(Telemetry)

monkey/monkey_island/cc/models/agent.py

@@ -4,8 +4,9 @@ from uuid import UUID
 
 from pydantic import Field
 
+from common.base_models import MutableBaseModel
+
 from . import MachineID
-from .base_models import MutableBaseModel
 
 
 class Agent(MutableBaseModel):

monkey/monkey_island/cc/models/machine.py

@@ -4,16 +4,16 @@ from typing import Optional, Sequence
 from pydantic import Field, PositiveInt, validator
 
 from common import OperatingSystem
-
-from .base_models import MutableBaseModel
-from .transforms import make_immutable_sequence
+from common.base_models import MutableBaseModel
+from common.transforms import make_immutable_sequence
+from common.types import HardwareID
 
 MachineID = PositiveInt
 
 
 class Machine(MutableBaseModel):
     id: MachineID = Field(..., allow_mutation=False)
-    hardware_id: Optional[PositiveInt]
+    hardware_id: Optional[HardwareID]
     network_interfaces: Sequence[IPv4Interface]
     operating_system: OperatingSystem
     operating_system_version: str

monkey/monkey_island/cc/models/node.py

@@ -2,9 +2,10 @@ from typing import Sequence, Tuple
 
 from pydantic import Field, validator
 
+from common.base_models import MutableBaseModel
+from common.transforms import make_immutable_nested_sequence
+
 from . import CommunicationType, MachineID
-from .base_models import MutableBaseModel
-from .transforms import make_immutable_nested_sequence
 
 ConnectionsSequence = Sequence[Tuple[MachineID, Sequence[CommunicationType]]]
 

monkey/monkey_island/cc/resources/__init__.py

@@ -9,3 +9,4 @@ from .agent_configuration import AgentConfiguration
 from .pba_file_upload import PBAFileUpload, LINUX_PBA_TYPE, WINDOWS_PBA_TYPE
 from .pba_file_download import PBAFileDownload
 from .events import Events
+from .agents import Agents

monkey/monkey_island/cc/resources/agents.py

@@ -0,0 +1,28 @@
+import json
+import logging
+from http import HTTPStatus
+
+from flask import make_response, request
+
+from common import AgentRegistrationData
+from monkey_island.cc.resources.AbstractResource import AbstractResource
+
+logger = logging.getLogger(__name__)
+
+
+class Agents(AbstractResource):
+    urls = ["/api/agents"]
+
+    def post(self):
+        try:
+            # Just parse for now
+            agent_registration_data = AgentRegistrationData(**request.json)
+
+            logger.debug(f"Agent registered: {agent_registration_data}")
+
+            return make_response({}, HTTPStatus.NO_CONTENT)
+        except (TypeError, ValueError, json.JSONDecodeError) as err:
+            return make_response(
+                {"error": f"Invalid configuration supplied: {err}"},
+                HTTPStatus.BAD_REQUEST,
+            )

monkey/tests/unit_tests/common/test_transforms.py

@@ -3,10 +3,7 @@ from typing import MutableSequence, Sequence
 
 import pytest
 
-from monkey_island.cc.models.transforms import (
-    make_immutable_nested_sequence,
-    make_immutable_sequence,
-)
+from common.transforms import make_immutable_nested_sequence, make_immutable_sequence
 
 
 def test_make_immutable_sequence__list():

monkey/tests/unit_tests/infection_monkey/utils/test_ids.py

@@ -0,0 +1,17 @@
+from uuid import UUID
+
+from infection_monkey.utils.ids import get_agent_id, get_machine_id
+
+
+def test_get_agent_id():
+    agent_id = get_agent_id()
+
+    assert isinstance(agent_id, UUID)
+    assert agent_id == get_agent_id()
+
+
+def test_get_machine_id():
+    machine_id = get_machine_id()
+
+    assert isinstance(machine_id, int)
+    assert machine_id == get_machine_id()

monkey/tests/unit_tests/monkey_island/cc/models/test_agent_registration_message.py

@@ -0,0 +1,103 @@
+from datetime import datetime, timezone
+from ipaddress import IPv4Interface
+from typing import MutableSequence, Sequence
+from uuid import UUID
+
+import pytest
+
+from common import AgentRegistrationData
+
+AGENT_ID = UUID("012e7238-7b81-4108-8c7f-0787bc3f3c10")
+PARENT_ID = UUID("0fc9afcb-1902-436b-bd5c-1ad194252484")
+
+AGENT_REGISTRATION_MESSAGE_OBJECT_DICT = {
+    "id": AGENT_ID,
+    "machine_hardware_id": 2,
+    "start_time": datetime.fromtimestamp(1660848408, tz=timezone.utc),
+    "parent_id": PARENT_ID,
+    "cc_server": "192.168.1.1:5000",
+    "network_interfaces": [IPv4Interface("10.0.0.1/24"), IPv4Interface("192.168.5.32/16")],
+}
+
+AGENT_REGISTRATION_MESSAGE_SIMPLE_DICT = {
+    "id": str(AGENT_ID),
+    "machine_hardware_id": 2,
+    "start_time": "2022-08-18T18:46:48+00:00",
+    "parent_id": str(PARENT_ID),
+    "cc_server": "192.168.1.1:5000",
+    "network_interfaces": ["10.0.0.1/24", "192.168.5.32/16"],
+}
+
+
+def test_to_dict():
+    a = AgentRegistrationData(**AGENT_REGISTRATION_MESSAGE_OBJECT_DICT)
+    simple_dict = AGENT_REGISTRATION_MESSAGE_SIMPLE_DICT.copy()
+
+    assert a.dict(simplify=True) == simple_dict
+
+
+def test_from_serialized():
+    from_serialized = AgentRegistrationData(**AGENT_REGISTRATION_MESSAGE_SIMPLE_DICT)
+    from_objects = AgentRegistrationData(**AGENT_REGISTRATION_MESSAGE_OBJECT_DICT)
+
+    assert from_serialized == from_objects
+
+
+@pytest.mark.parametrize(
+    "key, value",
+    [
+        ("id", 1),
+        ("machine_hardware_id", "not-an-int"),
+        ("start_time", None),
+        ("parent_id", 2.1),
+        ("cc_server", []),
+        ("network_interfaces", "not-a-list"),
+    ],
+)
+def test_construct_invalid_field__type_error(key, value):
+    invalid_type_dict = AGENT_REGISTRATION_MESSAGE_SIMPLE_DICT.copy()
+    invalid_type_dict[key] = value
+
+    with pytest.raises(TypeError):
+        AgentRegistrationData(**invalid_type_dict)
+
+
+@pytest.mark.parametrize(
+    "key, value",
+    [
+        ("machine_hardware_id", -1),
+        ("start_time", "not-a-date-time"),
+        ("network_interfaces", [1, "stuff", 3]),
+    ],
+)
+def test_construct_invalid_field__value_error(key, value):
+    invalid_value_dict = AGENT_REGISTRATION_MESSAGE_SIMPLE_DICT.copy()
+    invalid_value_dict[key] = value
+
+    with pytest.raises(ValueError):
+        AgentRegistrationData(**invalid_value_dict)
+
+
+@pytest.mark.parametrize(
+    "key, value",
+    [
+        ("id", PARENT_ID),
+        ("machine_hardware_id", 99),
+        ("start_time", 0),
+        ("parent_id", AGENT_ID),
+        ("cc_server", "10.0.0.1:4999"),
+        ("network_interfaces", ["10.0.0.1/24"]),
+    ],
+)
+def test_fields_immutable(key, value):
+    a = AgentRegistrationData(**AGENT_REGISTRATION_MESSAGE_OBJECT_DICT)
+
+    with pytest.raises(TypeError):
+        setattr(a, key, value)
+
+
+def test_network_interfaces_sequence_immutable():
+    a = AgentRegistrationData(**AGENT_REGISTRATION_MESSAGE_OBJECT_DICT)
+
+    assert isinstance(a.network_interfaces, Sequence)
+    assert not isinstance(a.network_interfaces, MutableSequence)

monkey/tests/unit_tests/monkey_island/cc/resources/test_agents.py

@@ -0,0 +1,42 @@
+from http import HTTPStatus
+from uuid import UUID
+
+from tests.unit_tests.monkey_island.conftest import get_url_for_resource
+
+from monkey_island.cc.resources import Agents
+
+AGENTS_URL = get_url_for_resource(Agents)
+
+AGENT_REGISTRATION_DICT = {
+    "id": UUID("6bfd8b64-43d8-4449-8c70-d898aca74ad8"),
+    "machine_hardware_id": 1,
+    "start_time": 0,
+    "parent_id": UUID("9d55ba33-95c2-417d-bd86-d3d11e47daeb"),
+    "cc_server": "10.0.0.1:5000",
+    "network_interfaces": ["10.1.1.2/24"],
+}
+
+
+def test_agent_registration(flask_client):
+    print(AGENTS_URL)
+    resp = flask_client.post(
+        AGENTS_URL,
+        json=AGENT_REGISTRATION_DICT,
+        follow_redirects=True,
+    )
+
+    assert resp.status_code == HTTPStatus.NO_CONTENT
+
+
+def test_agent_registration_invalid_data(flask_client):
+    agent_registration_dict = AGENT_REGISTRATION_DICT.copy()
+
+    agent_registration_dict["id"] = 1
+
+    resp = flask_client.post(
+        AGENTS_URL,
+        json=agent_registration_dict,
+        follow_redirects=True,
+    )
+
+    assert resp.status_code == HTTPStatus.BAD_REQUEST