forked from p15670423/monkey
UT: Replace test data with SocketAddress wherever required
This commit is contained in:
parent
0ae653fb72
commit
3b192a869e
|
@ -11,6 +11,7 @@ from common.agent_event_serializers import (
|
||||||
)
|
)
|
||||||
from common.agent_events import AbstractAgentEvent
|
from common.agent_events import AbstractAgentEvent
|
||||||
from common.agent_registration_data import AgentRegistrationData
|
from common.agent_registration_data import AgentRegistrationData
|
||||||
|
from common.types import SocketAddress
|
||||||
from infection_monkey.island_api_client import (
|
from infection_monkey.island_api_client import (
|
||||||
HTTPIslandAPIClient,
|
HTTPIslandAPIClient,
|
||||||
IslandAPIConnectionError,
|
IslandAPIConnectionError,
|
||||||
|
@ -20,7 +21,7 @@ from infection_monkey.island_api_client import (
|
||||||
IslandAPITimeoutError,
|
IslandAPITimeoutError,
|
||||||
)
|
)
|
||||||
|
|
||||||
SERVER = "1.1.1.1:9999"
|
SERVER = SocketAddress(ip="1.1.1.1", port="9999")
|
||||||
PBA_FILE = "dummy.pba"
|
PBA_FILE = "dummy.pba"
|
||||||
WINDOWS = "windows"
|
WINDOWS = "windows"
|
||||||
AGENT_ID = UUID("80988359-a1cd-42a2-9b47-5b94b37cd673")
|
AGENT_ID = UUID("80988359-a1cd-42a2-9b47-5b94b37cd673")
|
||||||
|
|
|
@ -7,6 +7,7 @@ from uuid import UUID
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from common import AgentRegistrationData
|
from common import AgentRegistrationData
|
||||||
|
from common.types import SocketAddress
|
||||||
from monkey_island.cc.island_event_handlers import handle_agent_registration
|
from monkey_island.cc.island_event_handlers import handle_agent_registration
|
||||||
from monkey_island.cc.models import Agent, CommunicationType, Machine
|
from monkey_island.cc.models import Agent, CommunicationType, Machine
|
||||||
from monkey_island.cc.repository import (
|
from monkey_island.cc.repository import (
|
||||||
|
@ -31,7 +32,7 @@ AGENT_REGISTRATION_DATA = AgentRegistrationData(
|
||||||
machine_hardware_id=MACHINE.hardware_id,
|
machine_hardware_id=MACHINE.hardware_id,
|
||||||
start_time=0,
|
start_time=0,
|
||||||
parent_id=None,
|
parent_id=None,
|
||||||
cc_server="192.168.1.1:5000",
|
cc_server=SocketAddress(ip="192.168.1.1", port="5000"),
|
||||||
network_interfaces=[IPv4Interface("192.168.1.2/24")],
|
network_interfaces=[IPv4Interface("192.168.1.2/24")],
|
||||||
)
|
)
|
||||||
|
|
||||||
|
@ -111,7 +112,7 @@ def test_existing_machine_updated__find_by_ip(handler, machine_repository):
|
||||||
machine_hardware_id=5,
|
machine_hardware_id=5,
|
||||||
start_time=0,
|
start_time=0,
|
||||||
parent_id=None,
|
parent_id=None,
|
||||||
cc_server="192.168.1.1:5000",
|
cc_server=SocketAddress(ip="192.168.1.1", port="5000"),
|
||||||
network_interfaces=[
|
network_interfaces=[
|
||||||
IPv4Interface("192.168.1.2/24"),
|
IPv4Interface("192.168.1.2/24"),
|
||||||
IPv4Interface("192.168.1.4/24"),
|
IPv4Interface("192.168.1.4/24"),
|
||||||
|
@ -215,7 +216,7 @@ def test_machine_interfaces_updated(handler, machine_repository):
|
||||||
machine_hardware_id=MACHINE.hardware_id,
|
machine_hardware_id=MACHINE.hardware_id,
|
||||||
start_time=0,
|
start_time=0,
|
||||||
parent_id=None,
|
parent_id=None,
|
||||||
cc_server="192.168.1.1:5000",
|
cc_server=SocketAddress(ip="192.168.1.1", port="5000"),
|
||||||
network_interfaces=[
|
network_interfaces=[
|
||||||
IPv4Interface("192.168.1.2/24"),
|
IPv4Interface("192.168.1.2/24"),
|
||||||
IPv4Interface("192.168.1.3/16"),
|
IPv4Interface("192.168.1.3/16"),
|
||||||
|
|
|
@ -6,16 +6,18 @@ from uuid import UUID
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from common import AgentRegistrationData
|
from common import AgentRegistrationData
|
||||||
|
from common.types import SocketAddress
|
||||||
|
|
||||||
AGENT_ID = UUID("012e7238-7b81-4108-8c7f-0787bc3f3c10")
|
AGENT_ID = UUID("012e7238-7b81-4108-8c7f-0787bc3f3c10")
|
||||||
PARENT_ID = UUID("0fc9afcb-1902-436b-bd5c-1ad194252484")
|
PARENT_ID = UUID("0fc9afcb-1902-436b-bd5c-1ad194252484")
|
||||||
|
SOCKET_ADDRESS = SocketAddress(ip="192.168.1.1", port="5000")
|
||||||
|
|
||||||
AGENT_REGISTRATION_MESSAGE_OBJECT_DICT = {
|
AGENT_REGISTRATION_MESSAGE_OBJECT_DICT = {
|
||||||
"id": AGENT_ID,
|
"id": AGENT_ID,
|
||||||
"machine_hardware_id": 2,
|
"machine_hardware_id": 2,
|
||||||
"start_time": datetime.fromtimestamp(1660848408, tz=timezone.utc),
|
"start_time": datetime.fromtimestamp(1660848408, tz=timezone.utc),
|
||||||
"parent_id": PARENT_ID,
|
"parent_id": PARENT_ID,
|
||||||
"cc_server": "192.168.1.1:5000",
|
"cc_server": SOCKET_ADDRESS,
|
||||||
"network_interfaces": [IPv4Interface("10.0.0.1/24"), IPv4Interface("192.168.5.32/16")],
|
"network_interfaces": [IPv4Interface("10.0.0.1/24"), IPv4Interface("192.168.5.32/16")],
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -24,7 +26,7 @@ AGENT_REGISTRATION_MESSAGE_SIMPLE_DICT = {
|
||||||
"machine_hardware_id": 2,
|
"machine_hardware_id": 2,
|
||||||
"start_time": "2022-08-18T18:46:48+00:00",
|
"start_time": "2022-08-18T18:46:48+00:00",
|
||||||
"parent_id": str(PARENT_ID),
|
"parent_id": str(PARENT_ID),
|
||||||
"cc_server": "192.168.1.1:5000",
|
"cc_server": SOCKET_ADDRESS.dict(simplify=True),
|
||||||
"network_interfaces": ["10.0.0.1/24", "192.168.5.32/16"],
|
"network_interfaces": ["10.0.0.1/24", "192.168.5.32/16"],
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -85,7 +87,7 @@ def test_construct_invalid_field__value_error(key, value):
|
||||||
("machine_hardware_id", 99),
|
("machine_hardware_id", 99),
|
||||||
("start_time", 0),
|
("start_time", 0),
|
||||||
("parent_id", AGENT_ID),
|
("parent_id", AGENT_ID),
|
||||||
("cc_server", "10.0.0.1:4999"),
|
("cc_server", SOCKET_ADDRESS),
|
||||||
("network_interfaces", ["10.0.0.1/24"]),
|
("network_interfaces", ["10.0.0.1/24"]),
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
|
@ -6,6 +6,7 @@ from uuid import uuid4
|
||||||
import mongomock
|
import mongomock
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
|
from common.types import SocketAddress
|
||||||
from monkey_island.cc.models import Agent
|
from monkey_island.cc.models import Agent
|
||||||
from monkey_island.cc.repository import (
|
from monkey_island.cc.repository import (
|
||||||
IAgentRepository,
|
IAgentRepository,
|
||||||
|
@ -54,6 +55,7 @@ AGENTS = (
|
||||||
*RUNNING_AGENTS,
|
*RUNNING_AGENTS,
|
||||||
*STOPPED_AGENTS,
|
*STOPPED_AGENTS,
|
||||||
)
|
)
|
||||||
|
CC_SERVER = SocketAddress(ip="127.0.0.1", port="1984")
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
|
@ -116,7 +118,7 @@ def test_upsert_agent__insert_empty_repository(empty_agent_repository):
|
||||||
def test_upsert_agent__update(agent_repository):
|
def test_upsert_agent__update(agent_repository):
|
||||||
agents = deepcopy(AGENTS)
|
agents = deepcopy(AGENTS)
|
||||||
agents[0].stop_time = datetime.now()
|
agents[0].stop_time = datetime.now()
|
||||||
agents[0].cc_server = "127.0.0.1:1984"
|
agents[0].cc_server = CC_SERVER
|
||||||
|
|
||||||
agent_repository.upsert_agent(agents[0])
|
agent_repository.upsert_agent(agents[0])
|
||||||
|
|
||||||
|
|
|
@ -16,7 +16,7 @@ AGENT_REGISTRATION_DICT = {
|
||||||
"machine_hardware_id": 1,
|
"machine_hardware_id": 1,
|
||||||
"start_time": 0,
|
"start_time": 0,
|
||||||
"parent_id": UUID("9d55ba33-95c2-417d-bd86-d3d11e47daeb"),
|
"parent_id": UUID("9d55ba33-95c2-417d-bd86-d3d11e47daeb"),
|
||||||
"cc_server": "10.0.0.1:5000",
|
"cc_server": {"ip": "10.0.0.1", "port": "5000"},
|
||||||
"network_interfaces": ["10.1.1.2/24"],
|
"network_interfaces": ["10.1.1.2/24"],
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue