Agent: Update type hints in control.py

.pre-commit-config.yaml

@@ -52,3 +52,10 @@ repos:
     rev: v2.3
     hooks:
       - id: vulture
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v0.971
+    hooks:
+      - id: mypy
+        additional_dependencies: [types-paramiko, types-python-dateutil, types-requests]
+        exclude: ^monkey/tests/
+        args: [--ignore-missing-imports]

monkey/infection_monkey/control.py

@@ -2,7 +2,7 @@ import json
 import logging
 import platform
 from socket import gethostname
-from typing import Mapping, Optional
+from typing import MutableMapping, Optional, Tuple
 
 import requests
 from requests.exceptions import ConnectionError
@@ -16,8 +16,6 @@ from infection_monkey.transport.tcp import TcpProxy
 from infection_monkey.utils import agent_process
 from infection_monkey.utils.environment import is_windows_os
 
-requests.packages.urllib3.disable_warnings()
-
 logger = logging.getLogger(__name__)
 
 PBA_FILE_DOWNLOAD = "https://%s/api/pba/download/%s"
@@ -29,11 +27,11 @@ class ControlClient:
     # https://github.com/guardicore/monkey/blob/133f7f5da131b481561141171827d1f9943f6aec/monkey/infection_monkey/telemetry/base_telem.py
     control_client_object = None
 
-    def __init__(self, server_address: str, proxies: Optional[Mapping[str, str]] = None):
+    def __init__(self, server_address: str, proxies: Optional[MutableMapping[str, str]] = None):
         self.proxies = {} if not proxies else proxies
         self.server_address = server_address
 
-    def wakeup(self, parent=None):
+    def wakeup(self, parent: str = None):
         if parent:
             logger.debug("parent: %s" % (parent,))
 
@@ -63,7 +61,7 @@ class ControlClient:
             timeout=MEDIUM_REQUEST_TIMEOUT,
         )
 
-    def find_server(self, default_tunnel=None):
+    def find_server(self, default_tunnel: str = None) -> bool:
         logger.debug(f"Trying to wake up with Monkey Island server: {self.server_address}")
         if default_tunnel:
             logger.debug("default_tunnel: %s" % (default_tunnel,))
@@ -95,7 +93,7 @@ class ControlClient:
                 logger.info("No tunnel found")
                 return False
 
-    def set_proxies(self, proxy_find):
+    def set_proxies(self, proxy_find: Tuple[str, str]):
         """
         Note: The proxy schema changes between different versions of requests and urllib3,
         which causes the machine to not open a tunnel back.
@@ -114,7 +112,7 @@ class ControlClient:
         else:
             self.proxies["https"] = f"{proxy_address}:{proxy_port}"
 
-    def send_telemetry(self, telem_category, json_data: str):
+    def send_telemetry(self, telem_category: str, json_data: str):
         if not self.server_address:
             logger.error(
                 "Trying to send %s telemetry before current server is established, aborting."
@@ -134,7 +132,7 @@ class ControlClient:
         except Exception as exc:
             logger.warning(f"Error connecting to control server {self.server_address}: {exc}")
 
-    def send_log(self, log):
+    def send_log(self, log: str):
         if not self.server_address:
             return
         try:
@@ -150,7 +148,7 @@ class ControlClient:
         except Exception as exc:
             logger.warning(f"Error connecting to control server {self.server_address}: {exc}")
 
-    def create_control_tunnel(self, keep_tunnel_open_time: int):
+    def create_control_tunnel(self, keep_tunnel_open_time: int) -> Optional[tunnel.MonkeyTunnel]:
         if not self.server_address:
             return None
 
@@ -158,8 +156,8 @@ class ControlClient:
         if my_proxy:
             proxy_class = TcpProxy
             try:
-                target_addr, target_port = my_proxy.split(":", 1)
+                target_addr, target_port_str = my_proxy.split(":", 1)
-                target_port = int(target_port)
+                target_port = int(target_port_str)
             except ValueError:
                 return None
         else:
@@ -173,7 +171,7 @@ class ControlClient:
             target_port=target_port,
         )
 
-    def get_pba_file(self, filename):
+    def get_pba_file(self, filename: str):
         try:
             return requests.get(  # noqa: DUO123
                 PBA_FILE_DOWNLOAD % (self.server_address, filename),

monkey/infection_monkey/utils/commands.py

@@ -1,3 +1,5 @@
+from typing import List
+
 from infection_monkey.config import GUID
 from infection_monkey.exploit.tools.helpers import AGENT_BINARY_PATH_LINUX, AGENT_BINARY_PATH_WIN64
 from infection_monkey.model import CMD_CARRY_OUT, CMD_EXE, MONKEY_ARG
@@ -9,7 +11,14 @@ DROPPER_TARGET_PATH_WIN64 = AGENT_BINARY_PATH_WIN64
 
 
 def build_monkey_commandline(target_host: VictimHost, depth: int, location: str = None) -> str:
+    """
+    Construct command line arguments from a VictimHost.
 
+    :param target_host: The host upon which the the new agent will run.
+    :param depth: The current network depth.
+    :param location: Path into which to copy the agent, defaults to None.
+    :return: A string containing the command line arguments
+    """
     return " " + " ".join(
         build_monkey_commandline_explicitly(
             GUID,
@@ -27,7 +36,7 @@ def build_monkey_commandline_explicitly(
     server: str = None,
     depth: int = None,
     location: str = None,
-) -> list:
+) -> List[str]:
     cmdline = []
 
     if parent is not None:
@@ -49,13 +58,28 @@ def build_monkey_commandline_explicitly(
     return cmdline
 
 
-def get_monkey_commandline_windows(destination_path: str, monkey_cmd_args: list) -> list:
+def get_monkey_commandline_windows(destination_path: str, monkey_cmd_args: List[str]) -> List[str]:
+    """
+    Build a command to run the agent on Windows.
+
+    :param destination_path: The path to the agent executable.
+    :param monkey_cmd_args: A list of command line arguments for the agent.
+    :return: The command, as a list of strings.
+    """
+
     monkey_cmdline = [CMD_EXE, CMD_CARRY_OUT, destination_path, MONKEY_ARG]
 
     return monkey_cmdline + monkey_cmd_args
 
 
-def get_monkey_commandline_linux(destination_path: str, monkey_cmd_args: list) -> list:
+def get_monkey_commandline_linux(destination_path: str, monkey_cmd_args: List[str]) -> List[str]:
+    """
+    Build a command to run the agent on Linux.
+
+    :param destination_path: The path to the agent executable.
+    :param monkey_cmd_args: A list of command line arguments for the agent.
+    :return: The command, as a list of strings.
+    """
     monkey_cmdline = [destination_path, MONKEY_ARG]
 
     return monkey_cmdline + monkey_cmd_args