forked from p15670423/monkey
UT: Implement test_handle_scan_data__upserts_machine
This commit is contained in:
parent
a0d6565c4a
commit
1c127781ca
|
@ -1,12 +1,15 @@
|
|||
from ipaddress import IPv4Address, IPv4Interface
|
||||
from itertools import count
|
||||
from unittest.mock import MagicMock
|
||||
from uuid import UUID
|
||||
|
||||
import pytest
|
||||
|
||||
from common import OperatingSystem
|
||||
from common.agent_events import PingScanEvent
|
||||
from common.types import PingScanData
|
||||
from common.types import PingScanData, SocketAddress
|
||||
from monkey_island.cc.agent_event_handlers import handle_scan_data
|
||||
from monkey_island.cc.models import Agent, CommunicationType, Machine
|
||||
from monkey_island.cc.repository import (
|
||||
IAgentRepository,
|
||||
IMachineRepository,
|
||||
|
@ -17,12 +20,27 @@ from monkey_island.cc.repository import (
|
|||
)
|
||||
|
||||
SEED_ID = 99
|
||||
AGENT_ID = UUID("1d8ce743-a0f4-45c5-96af-91106529d3e2")
|
||||
MACHINE_ID = 11
|
||||
CC_SERVER = SocketAddress(ip="10.10.10.100", port="5000")
|
||||
AGENT = Agent(id=AGENT_ID, machine_id=MACHINE_ID, start_time=0, parent_id=None, cc_server=CC_SERVER)
|
||||
MACHINE = Machine(
|
||||
id=MACHINE_ID,
|
||||
hardware_id=5,
|
||||
network_interfaces=[IPv4Interface("10.10.10.99/24")],
|
||||
)
|
||||
STORED_MACHINE = Machine(
|
||||
id=33,
|
||||
hardware_id=9,
|
||||
network_interfaces=[IPv4Interface("10.10.10.1/24")],
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def agent_repository() -> IAgentRepository:
|
||||
agent_repository = MagicMock(spec=IAgentRepository)
|
||||
agent_repository.upsert_agent = MagicMock()
|
||||
agent_repository.get_agent_by_id = MagicMock(return_value=AGENT)
|
||||
return agent_repository
|
||||
|
||||
|
||||
|
@ -42,12 +60,30 @@ def node_repository() -> INodeRepository:
|
|||
|
||||
|
||||
@pytest.fixture
|
||||
def handler() -> handle_scan_data:
|
||||
def handler(agent_repository, machine_repository, node_repository) -> handle_scan_data:
|
||||
return handle_scan_data(agent_repository, machine_repository, node_repository)
|
||||
|
||||
|
||||
def test_handle_scan_data__upserts_machine(handler: handle_scan_data):
|
||||
pass
|
||||
def test_handle_scan_data__upserts_machine(
|
||||
handler: handle_scan_data,
|
||||
machine_repository: IMachineRepository,
|
||||
node_repository: INodeRepository,
|
||||
):
|
||||
event = PingScanEvent(
|
||||
source=AGENT_ID,
|
||||
target=IPv4Address("10.10.10.1"),
|
||||
scan_data=PingScanData(True, OperatingSystem.LINUX),
|
||||
)
|
||||
machine_repository.get_machine_by_id = MagicMock(return_value=STORED_MACHINE)
|
||||
handler(event)
|
||||
|
||||
expected_machine = STORED_MACHINE.copy()
|
||||
expected_machine.operating_system = OperatingSystem.LINUX
|
||||
|
||||
assert machine_repository.upsert_machine.called_with(expected_machine)
|
||||
assert node_repository.upsert_communication.called_with(
|
||||
MACHINE.id, expected_machine.id, CommunicationType.SCANNED
|
||||
)
|
||||
|
||||
|
||||
def test_handle_scan_data__upserts_node(handler: handle_scan_data):
|
||||
|
|
Loading…
Reference in New Issue