Island: Add agent model

This commit is contained in:
Mike Salvatore 2022-08-18 15:07:06 -04:00
parent a625cc4583
commit 4a05c5a250
3 changed files with 152 additions and 0 deletions

View File

@ -12,3 +12,4 @@ from .user_credentials import UserCredentials
from .machine import Machine, MachineID from .machine import Machine, MachineID
from .communication_type import CommunicationType from .communication_type import CommunicationType
from .node import Node from .node import Node
from .agent import Agent

View File

@ -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="")

View File

@ -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