forked from p15670423/monkey
Merge branch 'i-node-repository' into develop
This commit is contained in:
commit
68bc73d86a
|
@ -12,4 +12,3 @@ class CommunicationType(Enum):
|
||||||
SCANNED = "scanned"
|
SCANNED = "scanned"
|
||||||
EXPLOITED = "exploited"
|
EXPLOITED = "exploited"
|
||||||
CC = "cc"
|
CC = "cc"
|
||||||
CC_TUNNEL = "cc_tunnel"
|
|
||||||
|
|
|
@ -1,19 +1,25 @@
|
||||||
from typing import Sequence, Tuple
|
from typing import Mapping, Tuple, TypeAlias
|
||||||
|
|
||||||
from pydantic import Field, validator
|
from pydantic import Field, validator
|
||||||
|
|
||||||
from common.base_models import MutableInfectionMonkeyBaseModel
|
from common.base_models import MutableInfectionMonkeyBaseModel
|
||||||
from common.transforms import make_immutable_nested_sequence
|
|
||||||
|
|
||||||
from . import CommunicationType, MachineID
|
from . import CommunicationType, MachineID
|
||||||
|
|
||||||
ConnectionsSequence = Sequence[Tuple[MachineID, Sequence[CommunicationType]]]
|
NodeConnections: TypeAlias = Mapping[MachineID, Tuple[CommunicationType, ...]]
|
||||||
|
|
||||||
|
|
||||||
class Node(MutableInfectionMonkeyBaseModel):
|
class Node(MutableInfectionMonkeyBaseModel):
|
||||||
machine_id: MachineID = Field(..., allow_mutation=False)
|
"""
|
||||||
connections: ConnectionsSequence
|
A network node and its outbound connections/communications
|
||||||
|
|
||||||
_make_immutable_nested_sequence = validator("connections", pre=True, allow_reuse=True)(
|
A node is identified by a MachineID and tracks all outbound communication to other machines on
|
||||||
make_immutable_nested_sequence
|
the network. This is particularly useful for creating graphs of Infection Monkey's activity
|
||||||
)
|
throughout the network.
|
||||||
|
"""
|
||||||
|
|
||||||
|
machine_id: MachineID = Field(..., allow_mutation=False)
|
||||||
|
"""The MachineID of the node (source)"""
|
||||||
|
|
||||||
|
connections: NodeConnections
|
||||||
|
"""All outbound connections from this node to other machines"""
|
||||||
|
|
|
@ -9,6 +9,7 @@ from .i_credentials_repository import ICredentialsRepository
|
||||||
from .i_user_repository import IUserRepository
|
from .i_user_repository import IUserRepository
|
||||||
from .i_machine_repository import IMachineRepository
|
from .i_machine_repository import IMachineRepository
|
||||||
from .i_agent_repository import IAgentRepository
|
from .i_agent_repository import IAgentRepository
|
||||||
|
from .i_node_repository import INodeRepository
|
||||||
|
|
||||||
|
|
||||||
from .local_storage_file_repository import LocalStorageFileRepository
|
from .local_storage_file_repository import LocalStorageFileRepository
|
||||||
|
|
|
@ -1,11 +0,0 @@
|
||||||
from abc import ABC
|
|
||||||
|
|
||||||
|
|
||||||
class INetworkMapRepository(ABC):
|
|
||||||
|
|
||||||
# TODO Define NetMap object
|
|
||||||
def get_map(self) -> NetMap: # noqa: F821
|
|
||||||
pass
|
|
||||||
|
|
||||||
def save_netmap(self, netmap: NetMap): # noqa: F821
|
|
||||||
pass
|
|
|
@ -0,0 +1,35 @@
|
||||||
|
from abc import ABC, abstractmethod
|
||||||
|
from typing import Sequence
|
||||||
|
|
||||||
|
from monkey_island.cc.models import CommunicationType, MachineID, Node
|
||||||
|
|
||||||
|
|
||||||
|
class INodeRepository(ABC):
|
||||||
|
"""A repository used to store and retrieve `Node` objects"""
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
|
def upsert_communication(
|
||||||
|
self, src: MachineID, dst: MachineID, communication_type: CommunicationType
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Insert or update a node connection
|
||||||
|
|
||||||
|
Insert or update data about how network nodes are able to communicate. Nodes are identified
|
||||||
|
by MachineID and store information about all outbound connections to other machines. By
|
||||||
|
providing a source machine, target machine, and how they communicated, nodes in this
|
||||||
|
repository can be created if they don't exist or updated if they do.
|
||||||
|
|
||||||
|
:param src: The machine that the connection or communication originated from
|
||||||
|
:param dst: The machine that the src communicated with
|
||||||
|
:param communication_type: The way the machines communicated
|
||||||
|
:raises StorageError: If an error occurred while attempting to upsert the Node
|
||||||
|
"""
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
|
def get_nodes(self) -> Sequence[Node]:
|
||||||
|
"""
|
||||||
|
Return all nodes that are stored in the repository
|
||||||
|
|
||||||
|
:return: All known Nodes
|
||||||
|
:raises RetrievalError: If an error occurred while attempting to retrieve the nodes
|
||||||
|
"""
|
|
@ -1,16 +1,16 @@
|
||||||
from typing import MutableSequence
|
from typing import MutableMapping, MutableSequence
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from monkey_island.cc.models import CommunicationType, Node
|
from monkey_island.cc.models import CommunicationType, MachineID, Node
|
||||||
|
|
||||||
|
|
||||||
def test_constructor():
|
def test_constructor():
|
||||||
machine_id = 1
|
machine_id = 1
|
||||||
connections = (
|
connections = {
|
||||||
(6, (CommunicationType.SCANNED,)),
|
6: (CommunicationType.SCANNED,),
|
||||||
(7, (CommunicationType.SCANNED, CommunicationType.EXPLOITED)),
|
7: (CommunicationType.SCANNED, CommunicationType.EXPLOITED),
|
||||||
)
|
}
|
||||||
n = Node(
|
n = Node(
|
||||||
machine_id=1,
|
machine_id=1,
|
||||||
connections=connections,
|
connections=connections,
|
||||||
|
@ -23,13 +23,10 @@ def test_constructor():
|
||||||
def test_serialization():
|
def test_serialization():
|
||||||
node_dict = {
|
node_dict = {
|
||||||
"machine_id": 1,
|
"machine_id": 1,
|
||||||
"connections": [
|
"connections": {
|
||||||
[
|
"6": ["cc", "scanned"],
|
||||||
6,
|
"7": ["exploited", "cc"],
|
||||||
["cc", "scanned"],
|
},
|
||||||
],
|
|
||||||
[7, ["exploited", "cc_tunnel"]],
|
|
||||||
],
|
|
||||||
}
|
}
|
||||||
n = Node(**node_dict)
|
n = Node(**node_dict)
|
||||||
|
|
||||||
|
@ -37,7 +34,7 @@ def test_serialization():
|
||||||
|
|
||||||
|
|
||||||
def test_machine_id_immutable():
|
def test_machine_id_immutable():
|
||||||
n = Node(machine_id=1, connections=[])
|
n = Node(machine_id=1, connections={})
|
||||||
|
|
||||||
with pytest.raises(TypeError):
|
with pytest.raises(TypeError):
|
||||||
n.machine_id = 2
|
n.machine_id = 2
|
||||||
|
@ -45,52 +42,45 @@ def test_machine_id_immutable():
|
||||||
|
|
||||||
def test_machine_id__invalid_type():
|
def test_machine_id__invalid_type():
|
||||||
with pytest.raises(TypeError):
|
with pytest.raises(TypeError):
|
||||||
Node(machine_id=None, connections=[])
|
Node(machine_id=None, connections={})
|
||||||
|
|
||||||
|
|
||||||
def test_machine_id__invalid_value():
|
def test_machine_id__invalid_value():
|
||||||
with pytest.raises(ValueError):
|
with pytest.raises(ValueError):
|
||||||
Node(machine_id=-5, connections=[])
|
Node(machine_id=-5, connections={})
|
||||||
|
|
||||||
|
|
||||||
def test_connections__mutable():
|
def test_connections__mutable():
|
||||||
n = Node(machine_id=1, connections=[])
|
n = Node(machine_id=1, connections={})
|
||||||
|
|
||||||
# Raises exception on failure
|
# Raises exception on failure
|
||||||
n.connections = [(5, []), (7, [])]
|
n.connections = {5: [], 7: []}
|
||||||
|
|
||||||
|
|
||||||
def test_connections__invalid_machine_id():
|
def test_connections__invalid_machine_id():
|
||||||
n = Node(machine_id=1, connections=[])
|
n = Node(machine_id=1, connections={})
|
||||||
|
|
||||||
with pytest.raises(ValueError):
|
with pytest.raises(ValueError):
|
||||||
n.connections = [(5, []), (-5, [])]
|
n.connections = {5: [], -5: []}
|
||||||
|
|
||||||
|
|
||||||
def test_connections__recursively_immutable():
|
def test_connections__recursively_immutable():
|
||||||
n = Node(
|
n = Node(
|
||||||
machine_id=1,
|
machine_id=1,
|
||||||
connections=[
|
connections={
|
||||||
[6, [CommunicationType.SCANNED]],
|
6: [CommunicationType.SCANNED],
|
||||||
[7, [CommunicationType.SCANNED, CommunicationType.EXPLOITED]],
|
7: [CommunicationType.SCANNED, CommunicationType.EXPLOITED],
|
||||||
],
|
},
|
||||||
)
|
)
|
||||||
|
|
||||||
assert not isinstance(n.connections, MutableSequence)
|
for connections in n.connections.values():
|
||||||
assert not isinstance(n.connections[0], MutableSequence)
|
assert not isinstance(connections, MutableSequence)
|
||||||
assert not isinstance(n.connections[1], MutableSequence)
|
|
||||||
assert not isinstance(n.connections[0][1], MutableSequence)
|
|
||||||
assert not isinstance(n.connections[1][1], MutableSequence)
|
|
||||||
|
|
||||||
|
|
||||||
def test_connections__set_invalid_communications_type():
|
def test_connections__set_invalid_communications_type():
|
||||||
connections = (
|
connections = {8: [CommunicationType.SCANNED, "invalid_comm_type"]}
|
||||||
[
|
|
||||||
[8, [CommunicationType.SCANNED, "invalid_comm_type"]],
|
|
||||||
],
|
|
||||||
)
|
|
||||||
|
|
||||||
n = Node(machine_id=1, connections=[])
|
n = Node(machine_id=1, connections={})
|
||||||
|
|
||||||
with pytest.raises(ValueError):
|
with pytest.raises(ValueError):
|
||||||
n.connections = connections
|
n.connections = connections
|
||||||
|
|
|
@ -18,7 +18,6 @@ from monkey_island.cc.repository.i_attack_repository import IAttackRepository
|
||||||
from monkey_island.cc.repository.i_config_repository import IConfigRepository
|
from monkey_island.cc.repository.i_config_repository import IConfigRepository
|
||||||
from monkey_island.cc.repository.i_log_repository import ILogRepository
|
from monkey_island.cc.repository.i_log_repository import ILogRepository
|
||||||
from monkey_island.cc.repository.i_machine_repository import IMachineRepository
|
from monkey_island.cc.repository.i_machine_repository import IMachineRepository
|
||||||
from monkey_island.cc.repository.i_network_map_repository import INetworkMapRepository
|
|
||||||
from monkey_island.cc.repository.i_report_repository import IReportRepository
|
from monkey_island.cc.repository.i_report_repository import IReportRepository
|
||||||
from monkey_island.cc.repository.i_simulation_repository import ISimulationRepository
|
from monkey_island.cc.repository.i_simulation_repository import ISimulationRepository
|
||||||
from monkey_island.cc.repository.i_telemetry_repository import ITelemetryRepository
|
from monkey_island.cc.repository.i_telemetry_repository import ITelemetryRepository
|
||||||
|
@ -260,8 +259,10 @@ IMachineRepository.upsert_machine
|
||||||
IMachineRepository.get_machine_by_id
|
IMachineRepository.get_machine_by_id
|
||||||
IMachineRepository.get_machine_by_hardware_id
|
IMachineRepository.get_machine_by_hardware_id
|
||||||
IMachineRepository.get_machines_by_ip
|
IMachineRepository.get_machines_by_ip
|
||||||
INetworkMapRepository.get_map
|
INodeRepository
|
||||||
INetworkMapRepository.save_netmap
|
INodeRepository.upsert_communication
|
||||||
|
INodeRepository.communication_type
|
||||||
|
INodeRepository.get_nodes
|
||||||
IReportRepository
|
IReportRepository
|
||||||
ISimulationRepository.save_simulation
|
ISimulationRepository.save_simulation
|
||||||
ISimulationRepository.get_simulation
|
ISimulationRepository.get_simulation
|
||||||
|
|
Loading…
Reference in New Issue