From 284ec3d1197fa250c31d54d91b94cc5812fe322a Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 26 Sep 2022 17:46:06 +0530 Subject: [PATCH] Agent: Use SocketAddress type in handle_agent_registration._get_or_create_cc_machine() --- .../cc/island_event_handlers/handle_agent_registration.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/monkey/monkey_island/cc/island_event_handlers/handle_agent_registration.py b/monkey/monkey_island/cc/island_event_handlers/handle_agent_registration.py index 363335c5f..2994f03c3 100644 --- a/monkey/monkey_island/cc/island_event_handlers/handle_agent_registration.py +++ b/monkey/monkey_island/cc/island_event_handlers/handle_agent_registration.py @@ -3,7 +3,7 @@ from ipaddress import IPv4Address, IPv4Interface from typing import List, Optional from common import AgentRegistrationData -from common.network.network_utils import address_to_ip_port +from common.types import SocketAddress from monkey_island.cc.models import Agent, CommunicationType, Machine from monkey_island.cc.repository import ( IAgentRepository, @@ -116,8 +116,8 @@ class handle_agent_registration: src_machine.id, dst_machine.id, CommunicationType.CC ) - def _get_or_create_cc_machine(self, cc_server: str) -> Machine: - dst_ip = IPv4Address(address_to_ip_port(cc_server)[0]) + def _get_or_create_cc_machine(self, cc_server: SocketAddress) -> Machine: + dst_ip = cc_server.ip try: return self._machine_repository.get_machines_by_ip(dst_ip)[0]