Fixed RESTAPI (issue #507)

.gitignore

@@ -8,3 +8,4 @@ downloads/*
 LastTask*
 setup/xar*
 setup/bomutils/*
+.venv

empire

@@ -3,7 +3,7 @@
 import sqlite3, argparse, sys, argparse, logging, json, string
 import os, re, time, signal, copy, base64, pickle
 from flask import Flask, request, jsonify, make_response, abort, url_for
-from time import localtime, strftime
+from time import localtime, strftime, sleep
 from OpenSSL import SSL
 from Crypto.Random import random
 
@@ -140,11 +140,11 @@ def get_permanent_token(conn):
 #
 ####################################################################
 
-def start_restful_api(startEmpire=False, suppress=False, username=None, password=None, port=1337):
+def start_restful_api(empireMenu, suppress=False, username=None, password=None, port=1337):
     """
     Kick off the RESTful API with the given parameters.
 
-    startEmpire -   start a complete Empire instance in the backend as well
+    empireMenu  -   Main empire menu object
     suppress    -   suppress most console output
     username    -   optional username to use for the API, otherwise pulls from the empire.db config
     password    -   optional password to use for the API, otherwise pulls from the empire.db config
@@ -155,6 +155,8 @@ def start_restful_api(startEmpire=False, suppress=False, username=None, password
 
     conn = database_connect()
 
+    main = empireMenu
+
     global serverExitCommand
 
     # if a username/password were not supplied, use the creds stored in the db
@@ -171,28 +173,8 @@ def start_restful_api(startEmpire=False, suppress=False, username=None, password
     else:
         execute_db_query(conn, "UPDATE config SET api_password=?", password)
 
-
-    class Namespace:
-        """
-        Temporary namespace to create the followin base argument object.
-        """
-        def __init__(self, **kwargs):
-            self.__dict__.update(kwargs)
-
-    # instantiate an Empire instance in case we need to interact with stagers or listeners
-    args = Namespace(debug=None, listener=None, stager=None, stager_options=None, version=False)
-
     print ''
 
-    if startEmpire:
-        # if we want to start a full-running empire instance
-        print " * Starting a full Empire instance"
-        main = empire.MainMenu(args=args)
-    else:
-        # if we just want the RESTful API, i.e. no listener/etc. startup
-        main = empire.MainMenu(args=args, restAPI=True)
-
-
     print " * Starting Empire RESTful API on port: %s" %(port)
 
     # refresh the token for the RESTful API
@@ -211,7 +193,6 @@ def start_restful_api(startEmpire=False, suppress=False, username=None, password
         # suppress all stdout and don't initiate the main cmdloop
         sys.stdout = open(os.devnull, 'w')
 
-
     # validate API token before every request except for the login URI
     @app.before_request
     def check_token():
@@ -1195,7 +1176,7 @@ def start_restful_api(startEmpire=False, suppress=False, username=None, password
         if conn:
             conn.close()
 
-        if startEmpire:
+        if suppress:
             print " * Shutting down the Empire instance"
             main.shutdown()
 
@@ -1234,15 +1215,16 @@ def start_restful_api(startEmpire=False, suppress=False, username=None, password
         signal.signal(signal.SIGINT, signal.default_int_handler)
         sys.exit()
 
-
+    try:
-    signal.signal(signal.SIGINT, signal_handler)
+        signal.signal(signal.SIGINT, signal_handler)
+    except ValueError:
+        pass
 
     # wrap the Flask connection in SSL and start it
     context = ('./data/empire.pem', './data/empire.pem')
     app.run(host='0.0.0.0', port=int(port), ssl_context=context, threaded=True)
 
 
-
 if __name__ == '__main__':
 
     parser = argparse.ArgumentParser()
@@ -1258,7 +1240,7 @@ if __name__ == '__main__':
 
     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('--rest', action='store_true', help='Run Empire and the 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')
@@ -1275,20 +1257,30 @@ if __name__ == '__main__':
         print empire.VERSION
 
     elif args.rest:
-        # start just the RESTful API
+        # start an Empire instance and RESTful API
-        while serverExitCommand == 'restart':
+        main = empire.MainMenu(args=args)
-            try:
+        def thread_api(empireMenu):
-                start_restful_api(startEmpire=False, suppress=False, username=args.username, password=args.password, port=args.restport)
+            while serverExitCommand == 'restart':
-            except SystemExit as e:
+                try:
-                pass
+                    start_restful_api(empireMenu=empireMenu, suppress=False, username=args.username, password=args.password, port=args.restport)
+                except SystemExit as e:
+                    pass
+
+        thread = helpers.KThread(target=thread_api, args=(main,))
+        thread.daemon = True
+        thread.start()
+        sleep(2)
+        main.cmdloop()
 
     elif args.headless:
         # start an Empire instance and RESTful API and suppress output
+        main = empire.MainMenu(args=args)
         while serverExitCommand == 'restart':
             try:
-                start_restful_api(startEmpire=True, suppress=True, username=args.username, password=args.password, port=args.restport)
+                start_restful_api(empireMenu=main, suppress=True, username=args.username, password=args.password, port=args.restport)
             except SystemExit as e:
                 pass
+
     else:
         # normal execution
         main = empire.MainMenu(args=args)

lib/common/empire.py

@@ -60,7 +60,7 @@ class MainMenu(cmd.Cmd):
     The main class used by Empire to drive the 'main' menu
     displayed when Empire starts.
     """
-    def __init__(self, args=None, restAPI=False):
+    def __init__(self, args=None):
 
         cmd.Cmd.__init__(self)