Handling metasploit errors with retries

app/metasploit.py

@@ -204,13 +204,15 @@ class Metasploit():
 
         payload_name = kwargs.get("outfile", "babymetal.exe")
         payload_type = kwargs.get("payload", None)
+        retries = 3
         if payload_type is None:
             raise MetasploitError("Payload not defined")
         try:
-            self.start_exploit_stub_for_external_payload(payload_type, lhost=kwargs.get("lhost", None))
+            ip = socket.gethostbyname(self.attacker.get_ip())
+            self.start_exploit_stub_for_external_payload(payload_type, lhost=kwargs.get("lhost", ip))
             self.wait_for_session(2)
         except MetasploitError:
-
+            while retries:
                 self.attack_logger.vprint(
                     f"{CommandlineColors.OKCYAN}Create payload {payload_name} {CommandlineColors.ENDC}",
                     1)
@@ -221,7 +223,12 @@ class Metasploit():
                     1)
 
                 self.start_exploit_stub_for_external_payload(payload=payload_type, lhost=kwargs.get("lhost", None))
-            self.wait_for_session()
+                try:
+                    self.wait_for_session(100)
+                    break
+                except MetasploitError:
+                    retries -= 1
+                    print(f"Global metasploit retries: {retries}")
 
 ##########################################################################