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

Agent: Reimplement HTTPIslandAPIClient.get_pba_file() as a method

This commit is contained in:
Mike Salvatore 2022-09-19 19:53:54 -04:00
parent aa3c6c2f4d
commit 841183d8e7
3 changed files with 9 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
monkey/infection_monkey/control.py
View File

@ -9,7 +9,7 @@ from urllib3 import disable_warnings
from common.common_consts.timeouts import MEDIUM_REQUEST_TIMEOUT from common.common_consts.timeouts import MEDIUM_REQUEST_TIMEOUT
from common.network.network_utils import get_my_ip_addresses from common.network.network_utils import get_my_ip_addresses
from infection_monkey.config import GUID from infection_monkey.config import GUID
from infection_monkey.island_api_client import HTTPIslandAPIClient, IIslandAPIClient from infection_monkey.island_api_client import IIslandAPIClient
from infection_monkey.network.info import get_host_subnets from infection_monkey.network.info import get_host_subnets
from infection_monkey.utils import agent_process from infection_monkey.utils import agent_process
@ -84,6 +84,6 @@ class ControlClient:
def get_pba_file(self, filename): def get_pba_file(self, filename):
try: try:
HTTPIslandAPIClient.get_pba_file(self.server_address, filename) return self._island_api_client.get_pba_file(filename)
except Exception as exc: except Exception as exc:
logger.warning(f"Error connecting to control server {self.server_address}: {exc}") logger.warning(f"Error connecting to control server {self.server_address}: {exc}")

9
monkey/infection_monkey/island_api_client/http_island_api_client.py
View File

@ -49,11 +49,12 @@ class HTTPIslandAPIClient(IIslandAPIClient):
timeout=MEDIUM_REQUEST_TIMEOUT, timeout=MEDIUM_REQUEST_TIMEOUT,
) )
@staticmethod
@handle_island_errors @handle_island_errors
def get_pba_file(server_address: str, filename: str): def get_pba_file(self, filename: str):
return requests.get( # noqa: DUO123 response = requests.get( # noqa: DUO123
"https://%s/api/pba/download/%s" % (server_address, filename), f"https://{self._island_server}/api/pba/download/{filename}",
verify=False, verify=False,
timeout=LONG_REQUEST_TIMEOUT, timeout=LONG_REQUEST_TIMEOUT,
) )
return response.content

4
monkey/infection_monkey/post_breach/custom_pba/custom_pba.py
View File

@ -76,7 +76,7 @@ class CustomPBA(PBA):
pba_file_contents = self.control_client.get_pba_file(filename) pba_file_contents = self.control_client.get_pba_file(filename)
status = None status = None
if not pba_file_contents or not pba_file_contents.content: if not pba_file_contents:
logger.error("Island didn't respond with post breach file.") logger.error("Island didn't respond with post breach file.")
status = ScanStatus.SCANNED status = ScanStatus.SCANNED
@ -97,7 +97,7 @@ class CustomPBA(PBA):
try: try:
with open(os.path.join(dst_dir, filename), "wb") as written_PBA_file: with open(os.path.join(dst_dir, filename), "wb") as written_PBA_file:
written_PBA_file.write(pba_file_contents.content) written_PBA_file.write(pba_file_contents)
return True return True
except IOError as e: except IOError as e:
logger.error("Can not upload post breach file to target machine: %s" % e) logger.error("Can not upload post breach file to target machine: %s" % e)