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

Agent: Fix return type hint in CachingAgentBinaryRepository._download_binary_from_island()

This commit is contained in:
Shreya Malviya 2022-09-20 15:19:34 +05:30
parent 4605722874
commit d84550ba23
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
monkey/infection_monkey/exploit/caching_agent_binary_repository.py
View File

@ -2,6 +2,7 @@ import io
import logging import logging
import threading import threading
from functools import lru_cache from functools import lru_cache
from typing import Optional
from common import OperatingSystem from common import OperatingSystem
from infection_monkey.island_api_client import IIslandAPIClient from infection_monkey.island_api_client import IIslandAPIClient
@ -33,9 +34,10 @@ class CachingAgentBinaryRepository(IAgentBinaryRepository):
return io.BytesIO(self._download_binary_from_island(operating_system)) return io.BytesIO(self._download_binary_from_island(operating_system))
@lru_cache(maxsize=None) @lru_cache(maxsize=None)
def _download_binary_from_island(self, operating_system: OperatingSystem) -> bytes: def _download_binary_from_island(self, operating_system: OperatingSystem) -> Optional[bytes]:
os_name = operating_system.value os_name = operating_system.value
try: try:
return self._island_api_client.get_agent_binary(os_name) return self._island_api_client.get_agent_binary(os_name)
except Exception as exc: except Exception as exc:
logger.warning(f"Error connecting to control server: {exc}") logger.warning(f"Error connecting to control server: {exc}")
return None