From b539e12347515e34e2f12777946e3e8df1cb1b50 Mon Sep 17 00:00:00 2001 From: byt3bl33d3r Date: Sat, 20 May 2017 15:33:03 -0600 Subject: [PATCH] Fixed RESTAPI (issue #507) --- .gitignore | 1 + empire | 64 +++++++++++++++++++------------------------- lib/common/empire.py | 2 +- 3 files changed, 30 insertions(+), 37 deletions(-) diff --git a/.gitignore b/.gitignore index 3e42ffa..100455a 100644 --- a/.gitignore +++ b/.gitignore @@ -8,3 +8,4 @@ downloads/* LastTask* setup/xar* setup/bomutils/* +.venv \ No newline at end of file diff --git a/empire b/empire index 539c978..7c91226 100755 --- a/empire +++ b/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() - - signal.signal(signal.SIGINT, signal_handler) + try: + 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 - while serverExitCommand == 'restart': - try: - start_restful_api(startEmpire=False, suppress=False, username=args.username, password=args.password, port=args.restport) - except SystemExit as e: - pass + # start an Empire instance and RESTful API + main = empire.MainMenu(args=args) + def thread_api(empireMenu): + while serverExitCommand == 'restart': + try: + 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) diff --git a/lib/common/empire.py b/lib/common/empire.py index 1f34269..17741e1 100644 --- a/lib/common/empire.py +++ b/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)