From 37f129fe8198e90e7bff73001823b878113e978a Mon Sep 17 00:00:00 2001
From: Harmj0y <will@harmj0y.net>
Date: Mon, 2 May 2016 18:13:38 -0400
Subject: [PATCH] tightened up argparse validation

---
 changelog |  4 ++++
 empire    | 28 ++++++++++++++++++----------
 2 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/changelog b/changelog
index 291d3ee..ce5b816 100644
--- a/changelog
+++ b/changelog
@@ -1,3 +1,7 @@
+5/2/2016
+---------
+-tightened up argparse validation
+
 4/24/2016
 ---------
 -Added credentials/get_spn_tickets to request SPN tickets.
diff --git a/empire b/empire
index c81508d..971b461 100755
--- a/empire
+++ b/empire
@@ -1170,22 +1170,30 @@ def start_restful_api(startEmpire=False, suppress=False, username=None, password
 if __name__ == '__main__':
 
     parser = argparse.ArgumentParser()
-    parser.add_argument('--debug', nargs='?', const='1', help='Debug level for output (default of 1).')
 
-    parser.add_argument('-s', '--stager', nargs='?', const="list", help='Specify a stager to generate. Lists all stagers if none is specified.')
-    parser.add_argument('-o', '--stager-options', nargs='*', help="Supply options to set for a stager in OPTION=VALUE format. Lists options if nothing is specified.")
-    parser.add_argument('-l', '--listener', nargs='?', const="list", help='Display listener options. Displays all listeners if nothing is specified.')
-    parser.add_argument('-v', '--version', action='store_true', help='Display current Empire version.')
-    parser.add_argument('--rest', action='store_true', help='Run the Empire RESTful API.')
-    parser.add_argument('--restport', nargs='?', help='Port to run the Empire RESTful API on.')
-    parser.add_argument('--headless', action='store_true', help='Run Empire and the RESTful API headless without the usual interface.')
-    parser.add_argument('--username', nargs='?', help='Start the RESTful API with the specified username instead of pulling from empire.db')
-    parser.add_argument('--password', nargs='?', help='Start the RESTful API with the specified password instead of pulling from empire.db')
+    generalGroup = parser.add_argument_group('General Options')
+    generalGroup.add_argument('--debug', nargs='?', const='1', help='Debug level for output (default of 1, 2 for msg display).')
+    generalGroup.add_argument('-v', '--version', action='store_true', help='Display current Empire version.')
+
+    cliGroup = parser.add_argument_group('CLI Payload Options')
+    cliGroup.add_argument('-l', '--listener', nargs='?', const="list", help='Display listener options. Displays all listeners if nothing is specified.')
+    cliGroup.add_argument('-s', '--stager', nargs='?', const="list", help='Specify a stager to generate. Lists all stagers if none is specified.')
+    cliGroup.add_argument('-o', '--stager-options', nargs='*', help="Supply options to set for a stager in OPTION=VALUE format. Lists options if nothing is specified.")
+    
+    restGroup = parser.add_argument_group('RESTful API Options')
+    launchGroup = restGroup.add_mutually_exclusive_group()
+    launchGroup.add_argument('--rest', action='store_true', help='Run the Empire RESTful API.')
+    launchGroup.add_argument('--headless', action='store_true', help='Run Empire and the RESTful API headless without the usual interface.')
+    restGroup.add_argument('--restport', type=int, nargs=1, help='Port to run the Empire RESTful API on.')
+    restGroup.add_argument('--username', nargs=1, help='Start the RESTful API with the specified username instead of pulling from empire.db')
+    restGroup.add_argument('--password', nargs=1, help='Start the RESTful API with the specified password instead of pulling from empire.db')
 
     args = parser.parse_args()
 
     if not args.restport:
         args.restport = '1337'
+    else:
+        args.restport = args.restport[0]
 
     if args.version:
         print empire.VERSION