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

Agent: Reduce duplication due to error handling

This commit is contained in:
Kekoa Kaaikala 2022-09-19 20:51:05 +00:00
parent 42633c066f
commit 51ce974c43
1 changed files with 28 additions and 35 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

63
monkey/infection_monkey/master/control_channel.py
View File

@ -1,4 +1,5 @@
import logging import logging
from functools import wraps
from typing import Optional, Sequence from typing import Optional, Sequence
from uuid import UUID from uuid import UUID
@ -24,12 +25,30 @@ disable_warnings() # noqa: DUO131
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
def handle_island_api_errors(func):
@wraps(func)
def wrapper(*args, **kwargs):
try:
print(args)
func(*args, **kwargs)
except (
IslandAPIConnectionError,
IslandAPIRequestError,
IslandAPIRequestFailedError,
IslandAPITimeoutError,
) as e:
raise IslandCommunicationError(e)
return wrapper
class ControlChannel(IControlChannel): class ControlChannel(IControlChannel):
def __init__(self, server: str, agent_id: str, api_client: IIslandAPIClient): def __init__(self, server: str, agent_id: str, api_client: IIslandAPIClient):
self._agent_id = agent_id self._agent_id = agent_id
self._control_channel_server = server self._control_channel_server = server
self._island_api_client = api_client self._island_api_client = api_client
@handle_island_api_errors
def register_agent(self, parent: Optional[UUID] = None): def register_agent(self, parent: Optional[UUID] = None):
agent_registration_data = AgentRegistrationData( agent_registration_data = AgentRegistrationData(
id=get_agent_id(), id=get_agent_id(),
@ -41,47 +60,21 @@ class ControlChannel(IControlChannel):
network_interfaces=get_network_interfaces(), network_interfaces=get_network_interfaces(),
) )
try: self._island_api_client.register_agent(agent_registration_data)
self._island_api_client.register_agent(agent_registration_data)
except (IslandAPIConnectionError, IslandAPITimeoutError) as e:
raise IslandCommunicationError(e)
@handle_island_api_errors
def should_agent_stop(self) -> bool: def should_agent_stop(self) -> bool:
if not self._control_channel_server: if not self._control_channel_server:
logger.error("Agent should stop because it can't connect to the C&C server.") logger.error("Agent should stop because it can't connect to the C&C server.")
return True return True
try: return self._island_api_client.should_agent_stop(
return self._island_api_client.should_agent_stop( self._control_channel_server, self._agent_id
self._control_channel_server, self._agent_id )
)
except (
IslandAPIConnectionError,
IslandAPIRequestError,
IslandAPIRequestFailedError,
IslandAPITimeoutError,
) as e:
raise IslandCommunicationError(e)
@handle_island_api_errors
def get_config(self) -> AgentConfiguration: def get_config(self) -> AgentConfiguration:
try: return self._island_api_client.get_config(self._control_channel_server)
return self._island_api_client.get_config(self._control_channel_server)
except (
IslandAPIConnectionError,
IslandAPIRequestError,
IslandAPIRequestFailedError,
IslandAPITimeoutError,
) as e:
raise IslandCommunicationError(e)
@handle_island_api_errors
def get_credentials_for_propagation(self) -> Sequence[Credentials]: def get_credentials_for_propagation(self) -> Sequence[Credentials]:
try: return self._island_api_client.get_credentials_for_propagation(self._control_channel_server)
return self._island_api_client.get_credentials_for_propagation(
self._control_channel_server
)
except (
IslandAPIConnectionError,
IslandAPIRequestError,
IslandAPIRequestFailedError,
IslandAPITimeoutError,
) as e:
raise IslandCommunicationError(e)