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

Island: Handle missing node in add_tcp_connections

This commit is contained in:
vakarisz 2022-10-06 14:45:56 +03:00
parent b0ec035909
commit 3bc2e4876f
2 changed files with 10 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
monkey/monkey_island/cc/repository/mongo_node_repository.py
View File

@ -51,6 +51,9 @@ class MongoNodeRepository(INodeRepository):
def add_tcp_connections(self, machine_id: MachineID, tcp_connections: TCPConnections):
node = self._get_node_by_id(machine_id)
if node is None:
node = Node(machine_id=machine_id, connections={})
for target, connections in tcp_connections.items():
if target in node.tcp_connections:
node.tcp_connections[target] = tuple({*node.tcp_connections[target], *connections})

7
monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_node_repository.py
View File

@ -232,3 +232,10 @@ def test_upsert_tcp_connections__port_already_present(node_repository):
modified_node = [node for node in nodes if node.machine_id == 4][0]
assert set(modified_node.tcp_connections) == set(ALL_TCP_CONNECTIONS)
assert len(modified_node.tcp_connections) == len(ALL_TCP_CONNECTIONS)
def test_upsert_tcp_connections__node_missing(node_repository):
node_repository.add_tcp_connections(999, TCP_CONNECTION_PORT_80)
nodes = node_repository.get_nodes()
modified_node = [node for node in nodes if node.machine_id == 999][0]
assert set(modified_node.tcp_connections) == set(TCP_CONNECTION_PORT_80)