From 6e48ea897c17ba1c1f69803b44ff8bf1dab20dc6 Mon Sep 17 00:00:00 2001
From: Oleksii Shevchuk <alxchk@gmail.com>
Date: Sat, 11 Mar 2017 13:42:03 +0200
Subject: [PATCH] Pass server transport to gen arguments, if no arguments
 specified

---
 pupy/pupygen.py         | 15 ++++++++-------
 pupy/pupylib/PupyCmd.py | 28 ++++++++++++++++++----------
 2 files changed, 26 insertions(+), 17 deletions(-)

diff --git a/pupy/pupygen.py b/pupy/pupygen.py
index 3f8eac4d..013352f0 100755
--- a/pupy/pupygen.py
+++ b/pupy/pupygen.py
@@ -365,6 +365,13 @@ def get_parser(base_parser, default_arch='x86', default_os='windows', default_fo
     parser.add_argument('--debug-scriptlets', action='store_true', help="don't catch scriptlets exceptions on the client for debug purposes")
     parser.add_argument('--debug', action='store_true', help="build with the debug template (the payload open a console)")
     parser.add_argument('--workdir', help='Set Workdir (Default = current workdir)')
+    parser.add_argument(
+        'launcher', choices=[
+            x for x in launchers.iterkeys()
+        ], default='connect', nargs='?',
+        help="Choose a launcher. Launchers make payloads behave differently at startup."
+    )
+    parser.add_argument('launcher_args', nargs=argparse.REMAINDER, help="launcher options")
 
     return parser
 
@@ -382,6 +389,7 @@ def pupygen(args):
     l=launchers[args.launcher]()
     while True:
         try:
+            print "LAUNCHER ARGS: {}".format(args.launcher_args)
             l.parse_args(args.launcher_args)
         except LauncherError as e:
             if str(e).strip().endswith("--host is required") and not "--host" in args.launcher_args:
@@ -499,13 +507,6 @@ def pupygen(args):
 if __name__ == '__main__':
     Credentials.DEFAULT_ROLE = 'CLIENT'
     parser = get_parser(argparse.ArgumentParser)
-    parser.add_argument(
-        'launcher', choices=[
-            x for x in launchers.iterkeys()
-        ], default='auto_proxy',
-        help="Choose a launcher. Launchers make payloads behave differently at startup."
-    )
-    parser.add_argument('launcher_args', nargs=argparse.REMAINDER, help="launcher options")
     try:
         pupygen(parser.parse_args())
     except Exception, e:
diff --git a/pupy/pupylib/PupyCmd.py b/pupy/pupylib/PupyCmd.py
index 3e17c176..1f25c79b 100644
--- a/pupy/pupylib/PupyCmd.py
+++ b/pupy/pupylib/PupyCmd.py
@@ -687,8 +687,9 @@ class PupyCmd(cmd.Cmd):
             self.display_success('Stopping DNSCNC')
             self.dnscnc.stop()
 
-        self.display_success('Stopping listener')
-        self.pupsrv.server.close()
+        if self.pupsrv.server:
+            self.display_success('Stopping listener')
+            self.pupsrv.server.close()
 
         self.display_success('Restarting')
         os.execv(argv0, argv)
@@ -786,18 +787,24 @@ class PupyCmd(cmd.Cmd):
             default_arch=default_arch,
             default_format=default_format
         )
-        arg_parser.add_argument('-t', '--transport', default=self.pupsrv.transport, help='Transport')
-        arg_parser.add_argument('launcher', nargs='?', default='connect', help='Launcher')
-        arg_parser.add_argument(
-            'launcher_args', nargs=REMAINDER,
-            default=self.pupsrv.transport_kwargs, help='Transport args'
-        )
+
+        print "DEBUG: {} / {}/{}".format(self.pupsrv.transport,
+                                             self.pupsrv.transport_kwargs, type(self.pupsrv.transport_kwargs))
 
         try:
             args = arg_parser.parse_args(shlex.split(arg))
         except PupyModuleExit:
             return
 
+        print "DEBUG: After parse: {}".format(args.launcher_args)
+        if not args.launcher_args:
+            args.launcher_args = [
+                x for x in [
+                    '-t', self.pupsrv.transport,
+                    self.pupsrv.transport_kwargs
+                ] if x
+            ]
+
         if self.pupsrv.httpd:
             wwwroot = self.config.get_folder('wwwroot')
             if not args.output_dir:
@@ -1070,8 +1077,9 @@ class PupyCmd(cmd.Cmd):
             self.display_success('Stopping DNSCNC')
             self.dnscnc.stop()
 
-        self.display_success('Stopping listener')
-        self.pupsrv.server.close()
+        if self.pupsrv.server:
+            self.display_success('Stopping listener')
+            self.pupsrv.server.close()
 
         return True