From 474a2ea185e74a7a52c84b2c509c93f3a9778f78 Mon Sep 17 00:00:00 2001
From: Thorsten Sick <thorsten.sick@avast.com>
Date: Thu, 26 Aug 2021 15:44:27 +0200
Subject: [PATCH] Mypy fixes

---
 Makefile              | 5 ++++-
 app/config.py         | 4 ++--
 app/machinecontrol.py | 4 ++--
 app/metasploit.py     | 4 ++--
 app/pluginmanager.py  | 4 ++--
 requirements.txt      | 6 ++++++
 6 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/Makefile b/Makefile
index e8bf497..53014cd 100644
--- a/Makefile
+++ b/Makefile
@@ -16,4 +16,7 @@ shipit: test
 
 # More detailed pylint tests.
 pylint:
-	pylint --rcfile=pylint.rc  *.py app/*.py plugins/base/*.py
\ No newline at end of file
+	pylint --rcfile=pylint.rc  *.py app/*.py plugins/base/*.py
+
+mypy:
+	mypy app/
\ No newline at end of file
diff --git a/app/config.py b/app/config.py
index 0b934c0..6b19c36 100644
--- a/app/config.py
+++ b/app/config.py
@@ -178,12 +178,12 @@ class ExperimentConfig():
         for attacker in self.raw_config["attackers"]:
             self._attackers.append(MachineConfig(self.raw_config["attackers"][attacker]))
 
-    def targets(self) -> [MachineConfig]:
+    def targets(self) -> list[MachineConfig]:
         """ Return config for targets as MachineConfig objects """
 
         return self._targets
 
-    def attackers(self) -> [MachineConfig]:
+    def attackers(self) -> list[MachineConfig]:
         """ Return config for attackers as MachineConfig objects """
 
         return self._attackers
diff --git a/app/machinecontrol.py b/app/machinecontrol.py
index 811a13b..b5ac66b 100644
--- a/app/machinecontrol.py
+++ b/app/machinecontrol.py
@@ -212,7 +212,7 @@ class Machine():
             plugin.install()
             self.attack_logger.vprint(f"{CommandlineColors.OKGREEN}Installed sensor: {name}{CommandlineColors.ENDC}", 2)
 
-    def get_sensors(self) -> [SensorPlugin]:
+    def get_sensors(self) -> list[SensorPlugin]:
         """ Returns a list of running sensors """
         return self.sensors
 
@@ -307,7 +307,7 @@ class Machine():
             plugin.install(self.vm_manager)
             # self.vulnerabilities.append(plugin)
 
-    def get_vulnerabilities(self) -> [VulnerabilityPlugin]:
+    def get_vulnerabilities(self) -> list[VulnerabilityPlugin]:
         """ Returns a list of installed vulnerabilities """
         return self.vulnerabilities
 
diff --git a/app/metasploit.py b/app/metasploit.py
index 262b858..180dd26 100644
--- a/app/metasploit.py
+++ b/app/metasploit.py
@@ -147,7 +147,7 @@ class Metasploit():
             retries -= 1
         raise MetasploitError(f"Could not find session for {target.get_ip()} Name resolution worked: {name_resolution_worked}")
 
-    def meterpreter_execute(self, cmds: [str], session_number: int, delay=0) -> str:
+    def meterpreter_execute(self, cmds: list[str], session_number: int, delay=0) -> list[str]:
         """ Executes commands on the meterpreter, returns results read from shell
 
         @param cmds: commands to execute, a list
@@ -164,7 +164,7 @@ class Metasploit():
             res.append(shell.read())
         return res
 
-    def meterpreter_execute_on(self, cmds: [str], target, delay=0) -> str:
+    def meterpreter_execute_on(self, cmds: list[str], target, delay=0) -> list[str]:
         """ Executes commands on the meterpreter, returns results read from shell
 
         @param cmds: commands to execute, a list
diff --git a/app/pluginmanager.py b/app/pluginmanager.py
index aa9b562..0ca66bd 100644
--- a/app/pluginmanager.py
+++ b/app/pluginmanager.py
@@ -35,7 +35,7 @@ class PluginManager():
         self.base = "plugins/**/*.py"
         self.attack_logger = attack_logger
 
-    def get_plugins(self, subclass, name_filter=None) -> [BasePlugin]:
+    def get_plugins(self, subclass, name_filter=None) -> list[BasePlugin]:
         """ Returns a list plugins matching specified criteria
 
 
@@ -46,7 +46,7 @@ class PluginManager():
 
         res = []
 
-        def get_handlers(a_plugin) -> [subclass]:
+        def get_handlers(a_plugin):
             return a_plugin.produce()
 
         plugin_dirs = set()
diff --git a/requirements.txt b/requirements.txt
index 14e4c13..648e431 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -13,3 +13,9 @@ sphinxcontrib.asciinema==0.3.2
 paramiko==2.7.2
 pymetasploit3==1.0.3
 pylint
+
+# Mypy stuff
+mypy
+types-PyYAML
+types-requests
+types-simplejson
\ No newline at end of file