p41275903
/
monkey
forked from p15670423/monkey
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

UT: Implement test_handle_scan_data__upserts_machine

This commit is contained in:
Kekoa Kaaikala 2022-09-27 19:53:10 +00:00
parent a0d6565c4a
commit 1c127781ca
1 changed files with 40 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

44
monkey/tests/unit_tests/monkey_island/cc/agent_event_handlers/test_handle_scan_data.py
View File

@ -1,12 +1,15 @@
from ipaddress import IPv4Address, IPv4Interface
from itertools import count from itertools import count
from unittest.mock import MagicMock from unittest.mock import MagicMock
from uuid import UUID from uuid import UUID
import pytest import pytest
from common import OperatingSystem
from common.agent_events import PingScanEvent 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.agent_event_handlers import handle_scan_data
from monkey_island.cc.models import Agent, CommunicationType, Machine
from monkey_island.cc.repository import ( from monkey_island.cc.repository import (
IAgentRepository, IAgentRepository,
IMachineRepository, IMachineRepository,
@ -17,12 +20,27 @@ from monkey_island.cc.repository import (
) )
SEED_ID = 99 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 @pytest.fixture
def agent_repository() -> IAgentRepository: def agent_repository() -> IAgentRepository:
agent_repository = MagicMock(spec=IAgentRepository) agent_repository = MagicMock(spec=IAgentRepository)
agent_repository.upsert_agent = MagicMock() agent_repository.upsert_agent = MagicMock()
agent_repository.get_agent_by_id = MagicMock(return_value=AGENT)
return agent_repository return agent_repository
@ -42,12 +60,30 @@ def node_repository() -> INodeRepository:
@pytest.fixture @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) return handle_scan_data(agent_repository, machine_repository, node_repository)
def test_handle_scan_data__upserts_machine(handler: handle_scan_data): def test_handle_scan_data__upserts_machine(
pass 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): def test_handle_scan_data__upserts_node(handler: handle_scan_data):