p20319658
/
monkey
forked from p34709852/monkey
1
0
Fork 0
Code Issues Pull Requests 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 threading
from functools import lru_cache
from typing import Optional
from common import OperatingSystem
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))
@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
try:
return self._island_api_client.get_agent_binary(os_name)
except Exception as exc:
logger.warning(f"Error connecting to control server: {exc}")
return None