Agent: Implement IslandApiClient

This commit is contained in:
vakarisz 2022-09-16 18:03:54 +03:00 committed by Ilija Lazoroski
parent 0b0f039474
commit 8386a006d0
4 changed files with 76 additions and 0 deletions

View File

@ -1 +1,2 @@
from infection_monkey.transport.http import LockedHTTPServer from infection_monkey.transport.http import LockedHTTPServer
from .island_api_client import IslandApiClient

View File

@ -0,0 +1,39 @@
import logging
import requests
from common.common_consts.timeouts import MEDIUM_REQUEST_TIMEOUT
from infection_monkey.transport.island_api_client_errors import (
IslandAPIConnectionError,
IslandAPIError,
IslandAPITimeoutError,
)
logger = logging.getLogger(__name__)
class IslandApiClient:
"""
Represents Island API client
"""
def __init__(self, island_server: str):
"""
Tries to connect to the island.
:param island_server: String representing the island ip address and port
:raises IslandAPIError:
"""
try:
requests.get( # noqa: DUO123
f"https://{island_server}/api?action=is-up",
verify=False,
timeout=MEDIUM_REQUEST_TIMEOUT,
)
except requests.exceptions.ConnectionError as err:
raise IslandAPIConnectionError(err)
except TimeoutError as err:
raise IslandAPITimeoutError(err)
except Exception as err:
raise IslandAPIError(err)

View File

@ -0,0 +1,30 @@
import pytest
import requests
import requests_mock
from infection_monkey.transport import IslandApiClient
from infection_monkey.transport.island_api_client_errors import (
IslandAPIConnectionError,
IslandAPIError,
IslandAPITimeoutError,
)
SERVER = "1.1.1.1:9999"
ISLAND_URI = f"https://{SERVER}/api?action=is-up"
@pytest.mark.parametrize(
"actual_error, expected_error",
[
(requests.exceptions.ConnectionError, IslandAPIConnectionError),
(TimeoutError, IslandAPITimeoutError),
(Exception, IslandAPIError),
],
)
def test_island_api_client(actual_error, expected_error):
with requests_mock.Mocker() as m:
m.get(ISLAND_URI, exc=actual_error)
with pytest.raises(expected_error):
IslandApiClient(SERVER)

View File

@ -9,6 +9,8 @@ from common.agent_configuration.agent_sub_configurations import (
) )
from common.credentials import Credentials, LMHash, NTHash from common.credentials import Credentials, LMHash, NTHash
from infection_monkey.exploit.log4shell_utils.ldap_server import LDAPServerFactory from infection_monkey.exploit.log4shell_utils.ldap_server import LDAPServerFactory
from infection_monkey.transport import IslandApiClient
from infection_monkey.transport.island_api_client_errors import IslandAPIRequestFailedError
from monkey_island.cc.event_queue import IslandEventTopic, PyPubSubIslandEventQueue from monkey_island.cc.event_queue import IslandEventTopic, PyPubSubIslandEventQueue
from monkey_island.cc.models import Report from monkey_island.cc.models import Report
from monkey_island.cc.models.networkmap import Arc, NetworkMap from monkey_island.cc.models.networkmap import Arc, NetworkMap
@ -328,3 +330,7 @@ CC_TUNNEL
IslandEventTopic.AGENT_CONNECTED IslandEventTopic.AGENT_CONNECTED
IslandEventTopic.CLEAR_SIMULATION_DATA IslandEventTopic.CLEAR_SIMULATION_DATA
IslandEventTopic.RESET_AGENT_CONFIGURATION IslandEventTopic.RESET_AGENT_CONFIGURATION
# TODO: Remove after #2292 is closed
IslandApiClient
IslandAPIRequestFailedError