Starting Metasploit unit tests

app/metasploit.py

@@ -28,9 +28,9 @@ class Metasploit():
         :param kwargs: Relevant ones: uri, port, server, username
         """
 
-        self.password = password
-        self.attack_logger = attack_logger
-        self.username = kwargs.get("username", None)
+        self.password: str = password
+        self.attack_logger: AttackLog = attack_logger
+        self.username: str = kwargs.get("username", None)
         self.kwargs = kwargs
         self.client = None
 
@@ -65,12 +65,15 @@ class Metasploit():
         print(res)
         return res
 
+    def __msfrpcd_cmd__(self):
+        return f"killall msfrpcd; nohup msfrpcd -P {self.password} -U {self.username} -S &"
+
     def start_msfrpcd(self):
         """ Starts the msfrpcs on the attacker. Metasploit must alredy be installed there ! """
 
-        cmd = f"killall msfrpcd; nohup msfrpcd -P {self.password} -U {self.username} -S &"
+        # cmd = f"killall msfrpcd; nohup msfrpcd -P {self.password} -U {self.username} -S &"
 
-        self.attacker.remote_run(cmd, disown=True)
+        self.attacker.remote_run(self.__msfrpcd_cmd__(), disown=True)
         # print("msfrpcd started")
         # breakpoint()
         time.sleep(3)
@@ -97,7 +100,7 @@ class Metasploit():
                 self.start_msfrpcd()
                 time.sleep(sleeptime)
                 sleeptime += 5
-                print("Failed getting connection to msfrpcd. Retries left: {retries}")
+                print(f"Failed getting connection to msfrpcd. Retries left: {retries}")
             retries -= 1
 
         if self.client is None:

tests/test_metasploit.py

@@ -0,0 +1,65 @@
+import unittest
+from unittest.mock import patch
+from app.metasploit import Metasploit
+from app.attack_log import AttackLog
+from pymetasploit3.msfrpc import MsfRpcClient
+import requests
+from app.exceptions import ServerError
+import time
+
+# https://docs.python.org/3/library/unittest.html
+
+
+class FakeAttacker():
+
+    def __init__(self):
+        pass
+
+    def remote_run(self, cmd, disown):
+        pass
+
+    def get_ip(self):
+        return "66.55.44.33"
+
+
+class TestMetasploit(unittest.TestCase):
+
+    def setUp(self) -> None:
+        with patch.object(time, "sleep") as _:
+            self.attack_logger = AttackLog(0)
+
+    def test_basic_init(self):
+        with patch.object(time, "sleep") as _:
+            m = Metasploit("FooBar", self.attack_logger)
+        self.assertEqual(m.password, "FooBar")
+        self.assertEqual(m.attack_logger, self.attack_logger)
+
+    def test_msfrpcd_cmd(self):
+        attacker = FakeAttacker()
+        with patch.object(time, "sleep") as _:
+            m = Metasploit("FooBar", self.attack_logger, attacker=attacker, username="Pennywise")
+        self.assertEqual(m.__msfrpcd_cmd__(), "killall msfrpcd; nohup msfrpcd -P FooBar -U Pennywise -S &")
+
+    def test_get_client_simple(self):
+        attacker = FakeAttacker()
+        with patch.object(time, "sleep") as _:
+            m = Metasploit("FooBar", self.attack_logger, attacker=attacker, username="Pennywise")
+            m.client = "Foo"
+            self.assertEqual(m.get_client(), "Foo")
+
+    def test_get_client_success(self):
+        attacker = FakeAttacker()
+        with patch.object(time, "sleep") as _:
+            m = Metasploit("FooBar", self.attack_logger, attacker=attacker, username="Pennywise")
+            with patch.object(MsfRpcClient, "__init__", return_value=None) as mock_method:
+                m.get_client()
+        mock_method.assert_called_once_with("FooBar", attacker=attacker, username="Pennywise", server="66.55.44.33")
+
+    def test_get_client_retries(self):
+        attacker = FakeAttacker()
+        with patch.object(time, "sleep") as _:
+            m = Metasploit("FooBar", self.attack_logger, attacker=attacker, username="Pennywise")
+            with self.assertRaises(ServerError):
+                with patch.object(MsfRpcClient, "__init__", side_effect=requests.exceptions.ConnectionError()) as mock_method:
+                    m.get_client()
+            mock_method.assert_called_with("FooBar", attacker=attacker, username="Pennywise", server="66.55.44.33")