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

Agent: Use SocketAddress type in handle_agent_registration._get_or_create_cc_machine()

This commit is contained in:
Shreya Malviya 2022-09-26 17:46:06 +05:30
parent ac633a6e75
commit 284ec3d119
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
monkey/monkey_island/cc/island_event_handlers/handle_agent_registration.py
View File

@ -3,7 +3,7 @@ from ipaddress import IPv4Address, IPv4Interface
from typing import List, Optional from typing import List, Optional
from common import AgentRegistrationData 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.models import Agent, CommunicationType, Machine
from monkey_island.cc.repository import ( from monkey_island.cc.repository import (
IAgentRepository, IAgentRepository,
@ -116,8 +116,8 @@ class handle_agent_registration:
src_machine.id, dst_machine.id, CommunicationType.CC src_machine.id, dst_machine.id, CommunicationType.CC
) )
def _get_or_create_cc_machine(self, cc_server: str) -> Machine: def _get_or_create_cc_machine(self, cc_server: SocketAddress) -> Machine:
dst_ip = IPv4Address(address_to_ip_port(cc_server)[0]) dst_ip = cc_server.ip
try: try:
return self._machine_repository.get_machines_by_ip(dst_ip)[0] return self._machine_repository.get_machines_by_ip(dst_ip)[0]