ace

#!/usr/bin/env python3

import argparse
import atexit
import collections
import datetime
import fnmatch
import getpass
import io
import json
import logging
import os
import os.path
import pwd
import re
import shutil
import signal
import smtplib
import socket
import sqlite3
import sys
import tempfile
import time
import trace
import traceback
import unittest
import uuid

parser = argparse.ArgumentParser(description="Simple Alert Queue")
parser.add_argument('--saq-home', required=False, dest='saq_home', default=None,
    help="Sets the SAQ_HOME environment variable.")
parser.add_argument('-L', '--logging-config-path', required=False, dest='logging_config_path', default=None,
    help="Path to the logging configuration file.")
parser.add_argument('-c', '--config-path', required=False, dest='config_paths', action='append', default=[],
    help="ACE configuration files. $SAQ_HOME/etc/saq.default.ini is always loaded, additional override default settings.")
#parser.add_argument('--single-threaded', required=False, dest='single_threaded', default=False, action='store_true',
    #help="Force execution to take place in a single process and a single thread.  Useful for manual debugging.")
parser.add_argument('--log-level', required=False, dest='log_level', default=None,
    help="Change the root log level.")
parser.add_argument('-u', '--user-name', required=False, dest='user_name', default=None,
    help="The user name of the ACE user executing the command. This information is required for some commands.")
parser.add_argument('--start', required=False, dest='start', default=False, action='store_true',
    help="Start the specified service.  Blocks keyboard unless --daemon (-d) is used.")
parser.add_argument('--stop', required=False, dest='stop', default=False, action='store_true',
    help="Stop the specified service.  Only applies to services started with --daemon (-d).")
parser.add_argument('-d', '--daemon', required=False, dest='daemon', default=False, action='store_true',
    help="Run this process as a daemon in the background.")
parser.add_argument('-k', '--kill-daemon', required=False, dest='kill_daemon', default=False, action='store_true',
    help="Kill the currently processing process.")
parser.add_argument('--force-alerts', required=False, dest='force_alerts', default=False, action='store_true',
    help="Force all analysis to always generate an alert.")
parser.add_argument('--relative-dir', required=False, dest='relative_dir', default=None,
    help="Assume all storage paths are relative to the given directory.  Defaults to current work directory.")
parser.add_argument('-p', '--provide-decryption-password', required=False, action='store_true', dest='provide_decryption_password', default=False,
    help="Prompt for the decryption password so that encrypted archive files can be automatically analyzed.")
parser.add_argument('--trace', required=False, action='store_true', dest='trace', default=False,
    help="Enable execution tracing (debugging option).")
parser.add_argument('-D', required=False, action='store_true', dest='debug_on_error', default=False,
    help="Break into pdb if an unhanled exception is thrown or an assertion fails.")

subparsers = parser.add_subparsers(dest='cmd')

# utility functions
def disable_proxy():
    for proxy_setting in [ 'http_proxy', 'https_proxy', 'ftp_proxy' ]:
        if proxy_setting in os.environ:
            logging.warning("removing proxy setting {0}".format(proxy_setting))
            del os.environ[proxy_setting]

def recurse_analysis(analysis, level=0, current_tree=[]):
    """Used to generate a textual display of the analysis results."""
    if not analysis:
        return

    if analysis in current_tree:
        return

    current_tree.append(analysis)

    if level > 0 and len(analysis.observables) == 0 and len(analysis.tags) == 0 and analysis.summary is None:
        return

    display = '{}{}{}'.format('\t' * level, 
                              '<' + '!' * len(analysis.detections) + '> ' if analysis.detections else '', 
                              analysis.summary if analysis.summary is not None else str(analysis))
    if analysis.tags:
        display += ' [ {} ] '.format(', '.join([x.name for x in analysis.tags]))
    
    print(display)

    for observable in analysis.observables:
        display = '{} * {}{}:{}'.format('\t' * level, 
                                        '<' + '!' * len(observable.detections) + '> ' if observable.detections else '', 
                                         observable.type, 
                                         observable.value)
        if observable.time is not None:
            display += ' @ {0}'.format(observable.time)
        if observable.directives:
            display += ' {{ {} }} '.format(', '.join([x for x in observable.directives]))
        if observable.tags:
            display += ' [ {} ] '.format(', '.join([x.name for x in observable.tags]))
        print(display)

        for observable_analysis in observable.all_analysis:
            recurse_analysis(observable_analysis, level + 1, current_tree)

def display_analysis(root):
    recurse_analysis(root)

    #observables = set(root.all_observables)
    #if observables:
        #print("{} OBSERVATIONS".format(len(observables)))
        #for observable in observables:
            #print("* {} - {}".format(observable.type, observable.value))

    tags = set(root.all_tags)
    if tags:
        print("{} TAGS".format(len(tags)))
        for tag in tags:
            print('* {}'.format(tag))

    detections = root.all_detection_points
    if detections:
        print("{} DETECTIONS FOUND (marked with <!> above)".format(len(detections)))
        for detection in detections:
            print('* {}'.format(detection))

def kill_daemon(daemon_name):

    import psutil

    daemon_pid_path = os.path.join(saq.DATA_DIR, 'var', 'daemon', daemon_name)
    if not os.path.exists(daemon_pid_path):
        logging.warning("daemon file {} does not exist (process not running?)".format(daemon_pid_path))
        sys.exit(1)

    deamon_pid = None
    try:
        with open(daemon_pid_path, 'r') as fp:
            daemon_pid = int(fp.read())
    except Exception as e:
        logging.error("cannot read PID from {}: {}".format(daemon_pid_path, e))
        sys.exit(1)

    try:
        parent = psutil.Process(daemon_pid)

        logging.info("sending SIGTERM to {}".format(parent.pid))
        # this should gracefully allow the system to come to a stop
        parent.terminate()

        try:
            parent.wait(timeout=60)
            logging.info("system shut down")
        except Exception as e:
            logging.error("unable to terminate process {}: {}".format(parent.pid, e))

            # but if it doesn't work then we walk the tree kill all processes
            for child in parent.children(recursive=True):
                try:
                    child.kill()
                    logging.info("killed child process {}".format(child.pid))
                except Exception as e:
                    logging.error("unable to kill child process {}: {}".format(child, e))

            try:
                parent.kill()
            except Exception as e:
                logging.error("unable to kill process {}: {}".format(parent.pid, e))

    except Exception as e:
        logging.error("unable to stop process {}: {}".format(parent, e))

    try:
        os.remove(daemon_pid_path)
    except Exception as e:
        logging.error("unable to delete PID file {}: {}".format(daemon_pid_path, e))
        sys.exit(1)

    sys.exit(0)

def daemonize(daemon_name):
    # are we already running?
    daemon_dir = os.path.join(saq.DATA_DIR, 'var', 'daemon')
    if not os.path.exists(daemon_dir):
        try:
            os.makedirs(daemon_dir)
        except Exception as e:
            logging.error("unable to create {}: {}".format(daemon_dir, str(e)))
            sys.exit(1)

    daemon_pid_path = os.path.join(saq.DATA_DIR, 'var', 'daemon', daemon_name)
    if os.path.exists(daemon_pid_path):
        pid = None
        with open(daemon_pid_path, 'r') as fp:
            pid = int(fp.read())
        
        if pid > 0:
            try:
                os.kill(pid, 0)
            except OSError:
                logging.warning("deamon PID file {} exists, but the process is not running. Proceeding to start engine.")
            else:
                import psutil
                process = psutil.Process(pid)
                logging.error("deamon PID file {} exists and '{}' is running with the contained PID={}: use command with -k to kill existing daemon".format(daemon_pid_path, process.name(), pid))
                sys.exit(1)

    pid = None

    # http://code.activestate.com/recipes/278731-creating-a-daemon-the-python-way/
    try:
        pid = os.fork()
        #print("got pid {}".format(pid))
    except OSError as e:
        logging.fatal("{} ({})".format(e.strerror, e.errno))
        sys.exit(1)

    if pid == 0:
        os.setsid()

        try:
            pid = os.fork()
            #print("got pid {}".format(pid))
        except OSError as e:
            logging.fatal("{} ({})".format(e.strerror, e.errno))
            sys.exit(1)

        if pid > 0:
            # write the pid to a file
            with open(daemon_pid_path, 'w') as fp:
                fp.write(str(pid))

            print("started background process {}".format(pid))
            #print("pid {} exiting...".format(os.getpid()))
            os._exit(0)
    else:
        #print("pid {} exiting...".format(os.getpid()))
        os._exit(0)

    import resource
    maxfd = resource.getrlimit(resource.RLIMIT_NOFILE)[1]
    if (maxfd == resource.RLIM_INFINITY):
        maxfd = MAXFD

        for fd in range(0, maxfd):
            try:
                os.close(fd)
            except OSError:   # ERROR, fd wasn't open to begin with (ignored)
                pass

    if (hasattr(os, "devnull")):
        REDIRECT_TO = os.devnull
    else:
        REDIRECT_TO = "/dev/null"

    os.open(REDIRECT_TO, os.O_RDWR)
    os.dup2(0, 1)
    os.dup2(0, 2)

# ============================================================================
# export observables
#

# XXX not sure we need this anymore

def export_observables(args):
    """Exports observables, mappings and some alert context into a sqlite database."""
    from saq.database import get_db_connection

    if os.path.exists(args.output_file):
        try:
            os.remove(args.output_file)
        except Exception as e:
            logging.error("unable to delete {}: {}".format(args.output_file, e))
            sys.exit(1)

    output_db = sqlite3.connect(args.output_file)
    output_c = output_db.cursor()
    output_c.execute("""CREATE TABLE observables ( type text NOT NULL, value text NOT NULL)""")
    output_c.execute("""CREATE UNIQUE INDEX i_type_value ON observables( type, value )""")
    output_c.execute("""CREATE INDEX i_value ON observables( value )""")

    i = 0
    skip = 0
    with get_db_connection() as input_db:
        input_c = input_db.cursor()
        input_c.execute("""SELECT type, value FROM observables""")
        for _type, _value in input_c:
            try:
                _value = _value.decode('utf-8')
            except UnicodeDecodeError as e:
                skip += 1
                continue

            output_c.execute("""INSERT INTO observables ( type, value ) VALUES ( ?, ? )""", (_type, _value))
            i += 1

    output_db.commit()
    output_db.close()

    logging.info("exported {} observables (skipped {})".format(i, skip))
    sys.exit(0)

export_observables_parser = subparsers.add_parser('export-observables',
    help="Exports observables into a sqlite database.")
export_observables_parser.add_argument('output_file', 
    help="Path to the sqlite database to create. Existing files are overwritten.")
export_observables_parser.set_defaults(func=export_observables)
    

# ============================================================================
# ssdeep compilation
#

# XXX get rid of this garbage

def compile_ssdeep(args):
    """Compile the etc/ssdeep.json CRITS export into a usable ssdeep hash file."""
    import saq

    json_path = os.path.join(saq.SAQ_HOME, args.input)
    output_path = os.path.join(saq.SAQ_HOME, args.output)
    temp_output_path = '{}.tmp'.format(output_path)
    count = 0
    duplicate_check = set()

    if not os.path.exists(json_path):
        logging.error("missing {}".format(json_path))
        sys.exit(1)

    try:
        with open(temp_output_path, 'w') as temp_fp:
            temp_fp.write('ssdeep,1.1--blocksize:hash:hash,filename\n')
            with open(json_path, 'r') as fp:
                root = json.load(fp)
                for _object in root['objects']:
                    if _object['ssdeep'] in duplicate_check:
                        logging.warning("detected duplicate ssdeep hash {}".format(_object['ssdeep']))
                        continue

                    duplicate_check.add(_object['ssdeep'])
                    temp_fp.write('{},{}\n'.format(_object['ssdeep'], _object['id']))
                    count += 1
    except Exception as e:
        logging.error("unable to create {}: {}".format(temp_output_path), str(e))
        sys.exit(1)

    # now move it over for usage
    try:
        logging.debug("moving {0} to {1}".format(temp_output_path, output_path))
        shutil.move(temp_output_path, output_path)
    except Exception as e:
        logging.error("unable to move {0} to {1}: {2}".format(temp_output_path, output_path, str(e)))
        sys.exit(1)

    logging.info("processed {0} entries".format(count))
    sys.exit(0)

compile_ssdeep_parser = subparsers.add_parser('compile-ssdeep',
    help="Compiles ssdeep signatures exported from CRITS.")
compile_ssdeep_parser.add_argument('-i', '--input', required=False, default='etc/ssdeep.json', dest='input',
    help="The CRITS export of the ssdeep hashses in JSON format.")
compile_ssdeep_parser.add_argument('-o', '--output', required=False, default='etc/ssdeep_hashes', dest='output',
    help="The ssdeep hash file generated by the script.")
compile_ssdeep_parser.set_defaults(func=compile_ssdeep)


# ============================================================================
# command line correlation
#

def correlate(args):
    # initialize command line engine
    from saq import CONFIG
    from saq.analysis import RootAnalysis
    from saq.constants import F_FILE, F_SUSPECT_FILE, VALID_OBSERVABLE_TYPES, event_time_format, \
                              EVENT_GLOBAL_TAG_ADDED, \
                              EVENT_GLOBAL_ANALYSIS_ADDED, \
                              EVENT_GLOBAL_OBSERVABLE_ADDED, \
                              ANALYSIS_MODE_CLI, ANALYSIS_MODE_CORRELATION
    from saq.engine import Engine
    from saq.error import report_exception
    from saq.process_server import initialize_process_server
    from saq.util import parse_event_time

    initialize_process_server()

    engine = Engine(single_threaded_mode=args.single_threaded)
    engine.set_local()
    engine.disable_alerting()

    def _cleanup():
        nonlocal engine
        from saq.database import get_db_connection
        with get_db_connection() as db:
            c = db.cursor()
            c.execute("DELETE FROM workload WHERE exclusive_uuid = %s", (engine.exclusive_uuid,))
            c.execute("DELETE FROM delayed_analysis WHERE exclusive_uuid = %s", (engine.exclusive_uuid,))
            c.execute("DELETE FROM nodes WHERE id = %s", (saq.SAQ_NODE_ID,))
            db.commit()

    # when all is said and done we want to make sure we don't leave any work entries behind in the database
    atexit.register(_cleanup)

    # these might not make sense any more
    #def _tag_added_event(source, event_type, tag):
        #logging.info("tagged {} with {}".format(source, tag.name))

    #def _observable_added_event(source, event_type, observable):
        #logging.info("observed {} in {}".format(observable, source))

    #def _analysis_added_event(source, event_type, analysis):
        #logging.info("analyzed {} with {}".format(source, analysis))
                #root.add_event_listener(EVENT_GLOBAL_TAG_ADDED, _tag_added_event)
                #root.add_event_listener(EVENT_GLOBAL_OBSERVABLE_ADDED, _observable_added_event)
                #root.add_event_listener(EVENT_GLOBAL_ANALYSIS_ADDED, _analysis_added_event)

    if len(args.targets) % 2 != 0:
        logging.error("odd number of arguments (you need pairs of type and value)")
        sys.exit(1)

    targets = args.targets
    
    # did we specify targets from stdin?
    if args.from_stdin:
        for o_value in sys.stdin:
            # the type of observables coming in on stdin is also specified on the command line
            targets.append(args.stdin_type)
            targets.append(o_value.strip())

    reference_time = None
    if args.reference_time is not None:
        reference_time = parse_event_time(args.reference_time)

    if os.path.exists(args.storage_dir) and not args.load:
        # if the output directory is the default directory then just delete it
        # this has been what I've wanted to happen 100% of the time
        if args.storage_dir == 'ace.out':
            try:
                logging.warning("deleting existing output directory ace.out")
                shutil.rmtree('ace.out')
            except Exception as e:
                logging.error("unable to delete existing output directory ace.out: {}".format(e))
                sys.exit(1)
        else:
            logging.error("output directory {} already exists".format(args.storage_dir))
            sys.exit(1)

    roots = []

    root = RootAnalysis()
    # create a RootAnalysis to pass to the engine for analysis
    root.storage_dir = args.storage_dir

    if os.path.exists(args.storage_dir):
        logging.warning("storage directory {} already exists".format(args.storage_dir))
    else:
        # create the output directory
        try:
            root.initialize_storage()
        except Exception as e:
            logging.error("unable to create output directory {}: {}".format(args.storage_dir, e))
            sys.exit(1)

    if args.load:
        root.load()

        # we override whatever previous analysis mode it had
        root.analysis_mode = ANALYSIS_MODE_CLI if args.analysis_mode is None else args.analysis_mode

    else:
        # set all of the properties individually
        # XXX only require company_id in RootAnalysis
        if args.company_name:
            root.company_name = args.company_name

        root.tool = 'ACE - Command Line Analysis'
        root.tool_instance = socket.gethostname()
        root.alert_type = args.alert_type
        root.analysis_mode = ANALYSIS_MODE_CLI if args.analysis_mode is None else args.analysis_mode
        root.description = args.description if args.description else 'Command Line Correlation'
        root.event_time = datetime.datetime.now() if reference_time is None else reference_time
        if args.load_details:
            with open(args.load_details, 'r') as fp:
                root.details = json.load(fp)
        else:
            root.details = { 
                'local_user': pwd.getpwuid(os.getuid())[0],
                'local_user_uid': os.getuid(),
                'comment': args.comment
            }

    # create the list of observables to add to the alert for analysis
    index = 0
    while index < len(args.targets):
        o_type = args.targets[index]
        o_value = args.targets[index + 1]

        if o_type not in VALID_OBSERVABLE_TYPES:
            logging.error("invalid observable type {0}".format(o_type))
            sys.exit(1)

        # the root analysis we're currently working on (defaults to the main alert)
        current_root = root

        # are we creating a separate alert for each observable?
        if args.split:
            current_root = RootAnalysis()
            subdir = os.path.join(root.storage_dir, current_root.uuid[0:3])
            if not os.path.exists(subdir):
                try:
                    os.mkdir(subdir)
                except Exception as e:
                    logging.error("unable to create directory {}: {}".format(subdir, e))
                    sys.exit(1)

            current_root.storage_dir = os.path.join(subdir, current_root.uuid)
            try:
                current_root.initialize_storage()
            except Exception as e:
                logging.error("unable to create directory {}: {}".format(subdir, e))

            # XXX not sure we need this any more
            # we'll make a little symlink if we can to help analysts know which directory is what
            # it's ok if this fails
            try:
                os.symlink(os.path.join(current_root.uuid[0:3], current_root.uuid), os.path.join(root.storage_dir, str(o_value)))
            except Exception as e:
                logging.warning("unable to create symlink: {}".format(e))

            current_root.tool = root.tool
            current_root.tool_instance = root.tool_instance
            current_root.alert_type = root.alert_type
            current_root.analysis_mode = root.analysis_mode
            current_root.description = "{} - {}:{}".format(root.description, o_type, o_value)
            current_root.event_time = root.event_time
            current_root.details = root.details

        # if this is a file then we need to copy it over to the storage directory
        if o_type == F_FILE:
            dest_path = os.path.join(current_root.storage_dir, os.path.basename(o_value))
            try:
                logging.debug("copying {} to {}".format(o_value, dest_path))
                shutil.copy(o_value, dest_path)
            except Exception as e:
                logging.error("unable to copy {} to {} for analysis: {}".format(o_value, dest_path, e))
                sys.exit(1)

            o_value = os.path.basename(o_value)

        observable = current_root.add_observable(o_type, o_value, reference_time)
        for directive in args.directives:
            observable.add_directive(directive)

        index += 2

        if args.split:
            current_root.save()
            current_root.schedule(exclusive_uuid=exclusive_uuid)
            roots.append(current_root)

    # if we are not splitting up the alerts then we just have one alert to look at
    if not args.split:
        root.save()
        root.schedule(exclusive_uuid=engine.exclusive_uuid)
        roots.append(root)

    # allow the user to control what analysis modules run
    if args.disable_all:
        logging.warning("disabling all analysis modules...")
        for section in saq.CONFIG.sections():
            if not section.startswith('analysis_module_'):
                continue

            saq.CONFIG[section]['enabled'] = 'no'

    if args.disabled_modules:
        for section in saq.CONFIG.sections():
            if not section.startswith('analysis_module_'):
                continue

        for name_pattern in args.disabled_modules:
            if name_pattern in section[len('analysis_module_'):]:
                logging.warning("disabling {}".format(section))
                saq.CONFIG[section]['enabled'] = 'no'

    # enable by group
    if args.enable_module_group:
        for module_group in args.enable_module_group:
            group_section = 'module_group_{}'.format(module_group)
            if group_section not in saq.CONFIG:
                logging.error("module group {} does not exist".format(module_group))
                sys.exit(1)

            for module_name in saq.CONFIG[group_section].keys():
                logging.info("enabling {} by group {}".format(module_name, module_group))
                engine.enable_module(module_name, ANALYSIS_MODE_CLI)
                #saq.CONFIG[module_name]['enabled'] = 'yes'

        #saq.CONFIG['engine_cli_correlation']['module_groups'] = ','.join(args.enable_module_group)

    if args.enabled_modules:
        for name_pattern in args.enabled_modules:
            for section in saq.CONFIG.sections():
                if fnmatch.fnmatch(section[len('analysis_module_'):], name_pattern):
                #if name_pattern in section[len('analysis_module_'):]:
                    logging.info("enabling {}".format(section))
                    engine.enable_module(section, ANALYSIS_MODE_CLI)
                    #saq.CONFIG[section]['enabled'] = 'yes'
                    #saq.CONFIG['engine_cli_correlation'][section] = 'yes'

    try:
        if not args.single_threaded:
            engine.controlled_stop()
        engine.start()
        engine.wait()
    except KeyboardInterrupt:
        logging.warning("user interrupted correlation")
        engine.stop()
        engine.wait()

    for root in roots:
        # display the results
        root = RootAnalysis(storage_dir=root.storage_dir)
        root.load()
        display_analysis(root)

        if args.alert:
            if root.whitelisted:
                logging.info("{} was whitelisted".format(root))
                continue

            if saq.FORCED_ALERTS or root.has_detections():
                from ace_api import upload
                try:
                    logging.info("uploading {}".format(root))

                    # need to switch the mode to correlation
                    root.analysis_mode = ANALYSIS_MODE_CORRELATION
                    root.save()

                    remote_host = None # if left as None then the api call defaults it to ace_api.default_node
                    if args.remote_host is not None:
                        remote_host = args.remote_host

                    if args.remote_port is not None:
                        remote_host = '{}:{}'.format(remote_host, args.remote_port)

                    result = upload(root.uuid, root.storage_dir, remote_host=remote_host, ssl_verification=args.ssl_ca_path)

                except Exception as e:
                    logging.error("unable to upload {}: {}".format(root, e))

    sys.exit(0)

# analyze-files
correlate_parser = subparsers.add_parser('correlate',
    help="Analyze one or more observables or alerts.")
correlate_parser.add_argument('--single-threaded', required=False, dest='single_threaded', default=False, action='store_true',
    help="Force execution to take place in a single process and a single thread.  Useful for manual debugging.")
correlate_parser.add_argument('-D', '--disable-module', required=False, dest='disabled_modules', nargs='*',
    help="Zero or more module names (substring match) to disable.")
correlate_parser.add_argument('--disable-all', required=False, dest='disable_all', default=False, action='store_true',
    help="Disable all analysis modules (use -E switch to enable specific modules.")
correlate_parser.add_argument('-E', '--enable-module', required=False, dest='enabled_modules', nargs='*', 
    help="Zero or more module names (substring match) to enable.")
correlate_parser.add_argument('-G', '--enable-module-group', required=False, dest='enable_module_group', nargs='+',
    help="One or more module group names to enable.")
correlate_parser.add_argument('-m', '--analysis-mode', required=False, dest='analysis_mode', 
    help="The analysis mode to use for this analysis. Defaults to cli")
correlate_parser.add_argument('-d', '--storage-dir', required=False, dest='storage_dir', default='ace.out',
    help="Specify an output directory.  Defaults to ace.out.")
correlate_parser.add_argument('-t', '--reference-time', required=False, dest='reference_time', default=None,
    help="Specify a datetime in YYYY-MM-DD HH:MM:SS format that observables (of a temporal type) should be referenced from.")
correlate_parser.add_argument('--description', required=False, dest='description', default="ACE Manual Correlation",
    help="Supply a description.  This will be displayed as part of the alert if this correlation is later imported as an alert.")
correlate_parser.add_argument('--comment', required=False, dest='comment', default=None,
    help="Optional generic comment to add to the details of the alert.")
correlate_parser.add_argument('--add-directive', required=False, dest='directives', action='append', default=[],
    help="Adds the given directive to all observables specified.  This option can be used multiple times.")
correlate_parser.add_argument('--alert-type', required=False, dest='alert_type', default='cli_analysis',
    help="Optionally set the alert type.  Some analysis is only performed for alerts of a certain type.")
correlate_parser.add_argument('--company-name', required=False, dest='company_name', default=None,
    help="Optionally assign ownership of this analysis to a company.")
correlate_parser.add_argument('--alert', required=False, dest='alert', action='store_true', default=False,
    help="Insert the correlation as an alert if it contains a detection point.")
correlate_parser.add_argument('--remote-host', required=False, dest='remote_host', default=None,
    help="Specify the remote host of the ACE system (defaults to the engine_ace configuration values).")
correlate_parser.add_argument('--remote-port', required=False, dest='remote_port', default=None, type=int,
    help="Specify the remote port of the ACE system (defaults to the engine_ace configuration values).")
#correlate_parser.add_argument('--ssl-root-cert', required=False, dest='ssl_root_cert', default=None,
    #help="Specify the path to the SSL cert for the ACE system (defaults to the engine_ace configuration values).")
#correlate_parser.add_argument('--ssl-key', required=False, dest='ssl_key', default=None,
    #help="Specify the path to the SSL key for the ACE system (defaults to the engine_ace configuration values).")
correlate_parser.add_argument('--ssl-ca-path', required=False, dest='ssl_ca_path', default=None,
    help="Specify the path to the CA cert for the ACE system (defaults to the engine_ace configuration values).")
#correlate_parser.add_argument('--ssl-hostname', required=False, dest='ssl_hostname', default=None,
    #help="Specify the ssl hostname of the ACE system (defaults to the engine_ace configuration values).")
correlate_parser.add_argument('--split', required=False, dest='split', action='store_true', default=False,
    help="Split the observables up into individual analysis.")
correlate_parser.add_argument('--from-stdin', required=False, dest='from_stdin', action='store_true', default=False,
    help="Read observables from stanard input.  Defaults to treating them as file-type observables.")
correlate_parser.add_argument('--stdin-type', required=False, dest='stdin_type', default='file',
    help="Specify the observable type when reading observables from stdin. Defaults to file.")
#correlate_parser.add_argument('--skip-analysis', required=False, dest='skip_analysis', action='store_true', default=False,
    #help="Skip analyzing the alert. Useful if you just want to send a bunch of stuff to ACE for analysis.")
correlate_parser.add_argument('--load', required=False, dest='load', action='store_true', default=False,
    help="Instead of creating a new analysis, load the existing analysis stored at --storage-dir.")
correlate_parser.add_argument('--load-details', required=False, dest='load_details', default=None,
    help="Load the given JSON file as the details of the alert.")
correlate_parser.add_argument('targets', nargs="*",
    help="One or more pairs of indicator types and values.")
correlate_parser.set_defaults(func=correlate)

# ============================================================================
# correlation engine
#

def correlation_engine(args):
    from saq.engine import Engine
    from saq.error import report_exception
    from saq.process_server import initialize_process_server

    daemon_name = 'correlation-engine'
    disable_proxy()

    if args.stop:
        kill_daemon(daemon_name)
        sys.exit(0)

    elif args.start:
        if args.daemon:
            logging.info("starting daemon {}".format(daemon_name))
            daemonize(daemon_name)

        initialize_process_server()

        ace_engine = None

        # add the capability for a graceful shutdown
        def handle_signal(signum, frame):
            ace_engine.stop()

        signal.signal(signal.SIGTERM, handle_signal)
        signal.signal(signal.SIGINT, handle_signal)

        try:
            ace_engine = Engine()
            ace_engine.start()
            ace_engine.wait()
        except Exception as e:
            logging.error("uncaught exception: {}".format(e))
            report_exception()
    
    else:
        logging.error("You must specify either --start or --stop for this service.")
        sys.exit(1)

    sys.exit(0)

correlation_engine_parser = subparsers.add_parser('correlation-engine',
    help="Start/Stop the ACE correlation analysis engine.")
correlation_engine_parser.set_defaults(func=correlation_engine)

# ============================================================================
# Remediation System
#

def remediation_system(args):
    from saq.remediation import initialize_remediation_system_manager, start_remediation_system_manager, stop_remediation_system_manager, wait_remediation_system_manager

    daemon_name = 'remediation-system'

    if args.stop:
        kill_daemon(daemon_name)
        sys.exit(0)

    elif args.start:
        if args.daemon:
            logging.info(f"starting daemon {daemon_name}")
            daemonize(daemon_name)
        
        try:
            initialize_remediation_system_manager()

            def handle_signal(signum, frame):
                stop_remediation_system_manager(wait=False)

            signal.signal(signal.SIGTERM, handle_signal)
            signal.signal(signal.SIGINT, handle_signal)

            start_remediation_system_manager()
            wait_remediation_system_manager()

        except KeyboardInterrupt:
            stop_remediation_system_manager()

        except Exception as e:
            logging.error(f"uncaught exception {e}")
            sys.exit(1)

        sys.exit(0)

    else:
        logging.error("You must specify either --start or --stop for this service.")
        sys.exit(1)

remediation_system_parser = subparsers.add_parser('remediation-system',
    help="Starts/stops the remediation system, responsible for implementing remediation requests in the background.")
remediation_system_parser.set_defaults(func=remediation_system)

# ============================================================================
# Message Dispatch System
#

def message_system(args):
    from saq.messaging import initialize_message_system, start_message_system, stop_message_system, wait_message_system

    daemon_name = 'message-dispatch-system'

    if args.stop:
        kill_daemon(daemon_name)
        sys.exit(0)

    elif args.start:
        if args.daemon:
            logging.info(f"starting daemon {daemon_name}")
            daemonize(daemon_name)
        
        try:
            initialize_message_system()

            def handle_signal(signum, frame):
                stop_message_system(wait=False)

            signal.signal(signal.SIGTERM, handle_signal)
            signal.signal(signal.SIGINT, handle_signal)

            start_message_system()
            wait_message_system()

        except KeyboardInterrupt:
            stop_message_system()

        except Exception as e:
            logging.error(f"uncaught exception {e}")
            sys.exit(1)

        sys.exit(0)

    else:
        logging.error("You must specify either --start or --stop for this service.")
        sys.exit(1)

message_system_parser = subparsers.add_parser('message-system',
    help="Starts/stops the message dispatch system, responsible for dispatching messages to other systems.")
message_system_parser.set_defaults(func=message_system)


# ============================================================================
# start GUI (non-apache version)
#

def start_gui(args):
    import saq
    from app import create_app, db
    from app.models import User
    from flask_script import Manager, Shell

    app = create_app()
    # add the "do" template command
    app.jinja_env.add_extension('jinja2.ext.do')

    if args.print_uri_paths:
        for rule in app.url_map.iter_rules():
            print(rule)
        sys.exit(0)

    app.run(saq.CONFIG.get('gui', 'listen_address'), debug=True, port=saq.CONFIG.getint('gui', 'listen_port'), 
        ssl_context=(saq.CONFIG.get('gui', 'ssl_cert'), saq.CONFIG.get('gui', 'ssl_key')),
        use_reloader=True)

# start-gui
start_gui_parser = subparsers.add_parser('start-gui',
    help="Start the SAQ GUI.")
start_gui_parser.add_argument('args', nargs=argparse.REMAINDER,
    help="Parameters to pass to the GUI command shell.")
start_gui_parser.add_argument('--print-uri-paths', default=False, action='store_true',
    help="Print all of the availble URL paths and exit, without starting the GUI.")
start_gui_parser.set_defaults(func=start_gui)

# ============================================================================
# start API (non-apache version)
#

def start_api(args):
    import saq
    from api import create_app

    app = create_app(testing=True)
    from werkzeug.serving import run_simple
    from werkzeug.wsgi import DispatcherMiddleware
    from flask import Flask, url_for
    app.config['DEBUG'] = True
    app.config['APPLICATION_ROOT'] = '/api'
    application = DispatcherMiddleware(Flask('dummy_app'), {
        app.config['APPLICATION_ROOT']: app,
    })

    if args.print_uri_paths:
        for rule in app.url_map.iter_rules():
            print(rule)
        sys.exit(0)

    run_simple(saq.CONFIG.get('api', 'listen_address'), saq.CONFIG.getint('api', 'listen_port'), application,
               ssl_context=(saq.CONFIG.get('api', 'ssl_cert'), saq.CONFIG.get('api', 'ssl_key')),
               use_reloader=False)

# start-gui
start_api_parser = subparsers.add_parser('start-api',
    help="Start the ACE API server in DEBUG mode.")
start_api_parser.add_argument('args', nargs=argparse.REMAINDER,
    help="Parameters to pass to the API command shell.")
start_api_parser.add_argument('--print-uri-paths', default=False, action='store_true',
    help="Print all of the availble URL paths and exit, without starting the API.")
start_api_parser.set_defaults(func=start_api)


# ============================================================================
# start/stop the primary ACE engine
#

def ace(args):

    from saq.engine import Engine
    from saq.error import report_exception
    from saq.process_server import initialize_process_server

    daemon_name = 'ace'
    disable_proxy()

    if args.stop:
        kill_daemon(daemon_name)
        sys.exit(0)

    elif args.start:
        if args.daemon:
            logging.info("starting daemon {}".format(daemon_name))
            daemonize(daemon_name)

        initialize_process_server()

        try:
            engine = Engine()
            engine.start()
            engine.wait()
        except KeyboardInterrupt as e:
            logging.info("caught user interrupt")
            try:
                engine.stop()
                engine.wait()
            except Exception as e:
                logging.error("uncaught exception: {}".format(e))
        except Exception as e:
            logging.error("uncaught exception: {}".format(e))
            report_exception()
    else:
        logging.error("You must specify either --start or --stop for this service.")
        sys.exit(1)

    sys.exit(0)

# ace-engine
ace_engine_parser = subparsers.add_parser('ace-engine',
    help="Start/Stop the Analysis Correlation Engine")
ace_engine_parser.set_defaults(func=ace)

# ============================================================================
# start/stop the carbon black binary collector
#

def cb_binary_collector(args):

    from saq.collectors.cb_binaries import CarbonBlackBinaryCollector
    from saq.error import report_exception

    daemon_name = 'cb_binary'
    disable_proxy()

    if args.stop:
        kill_daemon(daemon_name)
        sys.exit(0)

    elif args.start:
        if args.daemon:
            logging.info("starting daemon {}".format(daemon_name))
            daemonize(daemon_name)

        try:
            kwargs = {}
            if args.download_batch_size is not None:
                kwargs['download_batch_size'] = args.download_batch_size
            if args.initial_search_offset is not None:
                kwargs['initial_search_offset'] = args.initial_search_offset
            if args.storage_dir is not None:
                kwargs['storage_dir'] = args.storage_dir
            if args.collection_frequency is not None:
                kwargs['collection_frequency'] = args.collection_frequency

            collector = CarbonBlackBinaryCollector(**kwargs)
            collector.load_groups()
            collector.start()

        except KeyboardInterrupt as e:
            logging.info("caught user interrupt")
            try:
                collector.stop()
                collector.wait()
            except Exception as e:
                logging.error("uncaught exception: {}".format(e))
        except Exception as e:
            logging.error("uncaught exception: {}".format(e))
            report_exception()
    else:
        logging.error("You must specify either --start or --stop for this service.")
        sys.exit(1)

    sys.exit(0)

# cb-binary-collector
cb_binary_collector_parser = subparsers.add_parser('cb-binary-collector',
    help="Start/Stop the Carbon Black Binary Collector")
cb_binary_collector_parser.add_argument('-c', '--collection-frequency', type=int, default=10,
    help="The number of seconds to wait between requests to Carbon Black for new binaries. Defaults to 10 seconds.")
cb_binary_collector_parser.add_argument('--download-batch-size', type=int, default=None,
    help="The total number of carbon black binaries to download at once.")
cb_binary_collector_parser.add_argument('--initial-search-offset', type=int, default=None,
    help="The initial offset to use (in hours) if running this for the first time.")
#cb_binary_collector_parser.add_argument('--search-offset', type=int, default=None,
    #help="The number of minutes to look back for new binaries.")
cb_binary_collector_parser.add_argument('--storage-dir',  default=None,
    help="The directory to store the binaries in. Relative directories are relative to DATA_DIR.")
cb_binary_collector_parser.set_defaults(func=cb_binary_collector)

# ============================================================================
# start/stop the BRO HTTP stream collector
#

def bro_http_collector(args):

    from saq.collectors.http import BroHTTPStreamCollector
    from saq.error import report_exception

    daemon_name = 'bro_http'
    disable_proxy()

    if args.stop:
        kill_daemon(daemon_name)
        sys.exit(0)

    elif args.start:
        if args.daemon:
            logging.info("starting daemon {}".format(daemon_name))
            daemonize(daemon_name)

        try:
            collector = BroHTTPStreamCollector()
            collector.load_groups()
            collector.start()

        except KeyboardInterrupt as e:
            logging.info("caught user interrupt")
            try:
                collector.stop()
                collector.wait()
            except Exception as e:
                logging.error("uncaught exception: {}".format(e))
        except Exception as e:
            logging.error("uncaught exception: {}".format(e))
            report_exception()
    else:
        logging.error("You must specify either --start or --stop for this service.")
        sys.exit(1)

    sys.exit(0)

bro_http_collector_parser = subparsers.add_parser('bro-http-collector',
    help="Start/Stop the Bro HTTP Stream Collector")
bro_http_collector_parser.set_defaults(func=bro_http_collector)

# ============================================================================
# start/stop the BRO SMTP stream collector
#

def bro_smtp_collector(args):

    from saq.collectors.smtp import BroSMTPStreamCollector
    from saq.error import report_exception

    daemon_name = 'bro_smtp'
    disable_proxy()

    if args.stop:
        kill_daemon(daemon_name)
        sys.exit(0)

    elif args.start:
        if args.daemon:
            logging.info("starting daemon {}".format(daemon_name))
            daemonize(daemon_name)

        try:
            collector = BroSMTPStreamCollector()
            collector.load_groups()
            collector.start()

        except KeyboardInterrupt as e:
            logging.info("caught user interrupt")
            try:
                collector.stop()
                collector.wait()
            except Exception as e:
                logging.error("uncaught exception: {}".format(e))
        except Exception as e:
            logging.error("uncaught exception: {}".format(e))
            report_exception()
    else:
        logging.error("You must specify either --start or --stop for this service.")
        sys.exit(1)

    sys.exit(0)

bro_smtp_collector_parser = subparsers.add_parser('bro-smtp-collector',
    help="Start/Stop the Bro SMTP Stream Collector")
bro_smtp_collector_parser.set_defaults(func=bro_smtp_collector)

# ============================================================================
# start/stop the email collector
#

def email_collector(args):

    from saq.collectors.email import EmailCollector
    from saq.error import report_exception

    daemon_name = 'email'
    disable_proxy()

    if args.stop:
        kill_daemon(daemon_name)
        sys.exit(0)

    elif args.start:
        if args.daemon:
            logging.info("starting daemon {}".format(daemon_name))
            daemonize(daemon_name)

        try:
            collector = EmailCollector(assignment_yara_rule_path=args.assignment_rule,
                                       blacklist_yara_rule_path=args.blacklist_rule)
            collector.load_groups()
            collector.start()

        except KeyboardInterrupt as e:
            logging.info("caught user interrupt")
            try:
                collector.stop()
                collector.wait()
            except Exception as e:
                logging.error("uncaught exception: {}".format(e))
        except Exception as e:
            logging.error("uncaught exception: {}".format(e))
            report_exception()
    else:
        logging.error("You must specify either --start or --stop for this service.")
        sys.exit(1)

    sys.exit(0)

email_collector_parser = subparsers.add_parser('email-collector',
    help="Start/Stop the Email Collector")
email_collector_parser.add_argument('--blacklist-rule', default='etc/blacklist.yar',
    help="Path to the yara rule used for blacklisting emails.")
email_collector_parser.add_argument('--assignment-rule', default='etc/remote_assignments.yar',
    help="Path to the yara rule used for assigning emails to collection groups.")
email_collector_parser.set_defaults(func=email_collector)

# ============================================================================
# start carbon black binary processor
#

def carbon_black(args):
    from saq.engine.cb_engine import CarbonBlackCollectionEngine
    from saq.error import report_exception
    from saq.process_server import initialize_process_server

    daemon_name = 'carbon-black'
    disable_proxy()

    if args.stop:
        kill_daemon(daemon_name)
        sys.exit(0)

    elif args.start:
        if args.daemon:
            logging.info("starting daemon {0}".format(daemon_name))
            daemonize(daemon_name)

        initialize_process_server()

        cb_engine = None

        # add the capability for a graceful shutdown
        def handle_sigterm(signum, frame):
            logging.warning("received SIGTERM")
            cb_engine.stop()

        signal.signal(signal.SIGTERM, handle_sigterm)
    
        try:
            cb_engine = CarbonBlackCollectionEngine()
            cb_engine.start()
            cb_engine.wait()
        except KeyboardInterrupt as e:
            logging.warning("caught user interrupt")
            cb_engine.wait()
        except Exception as e:
            logging.error("uncaught exception: {}".format(e))
            report_exception()
    else:
        logging.error("You must specify either --start or --stop for this service.")
        sys.exit(1)

    sys.exit(0)

# carbon-black-binary
carbon_black_parser = subparsers.add_parser('carbon-black',
    help="Start/Stop the Carbon Black binary analysis engine.")
carbon_black_parser.set_defaults(func=carbon_black)

# ============================================================================
# start Network Semaphore
#

def network_semaphore(args):
    import saq
    from saq.network_semaphore import NetworkSemaphoreServer
    daemon_name = 'network-semaphore'

    if args.stop:
        kill_daemon(daemon_name)
        sys.exit(0)

    elif args.start:
        if args.daemon:
            logging.info("starting daemon {}".format(daemon_name))
            daemonize(daemon_name)

        server = NetworkSemaphoreServer()
        server.start()

        # add the capability for a graceful shutdown
        def handle_sigterm(signum, frame):
            logging.warning("received SIGTERM")
            server.stop()

        try:
            while True:
                #logging.debug("process {} sleeping".format(os.getpid()))
                time.sleep(1)
        except KeyboardInterrupt:
            server.stop()
    else:
        logging.error("You must specify either --start or --stop for this service.")
        sys.exit(1)

    sys.exit(0)

# network semaphore
network_semaphore_parser = subparsers.add_parser('network-semaphore',
    help="Start/Stop the Network Semaphore Server on this server.")
network_semaphore_parser.set_defaults(func=network_semaphore)

# ============================================================================
# user management
#

def add_user(args):
    from app.models import User
    from saq.database import get_db_connection
    from getpass import getpass


    u = User()
    u.username = args.username
    u.email = args.email

    try:
        import pytz
        u.timezone = 'Etc/UTC'
        if args.timezone:
            u.timezone = args.timezone

        pytz.timezone(u.timezone)

    except Exception as e:
        print(f"ERROR: invalid timezone {u.timezone}")
        sys.exit(1)

    password = getpass("Enter password for {0}: ".format(u.username))
    confirm = getpass("Confirm password for {0}: ".format(u.username))
    if password != confirm:
        logging.error("passwords do not match")
        sys.exit(1)

    u.password = password

    with get_db_connection() as db:
        c = db.cursor()
        c.execute("""INSERT INTO users ( username, email, password_hash, timezone ) VALUES ( %s, %s, %s, %s )""", (
            u.username, u.email, u.password_hash, u.timezone))
        db.commit()

    logging.info("added user {0}".format(u.username))

add_user_parser = subparsers.add_parser('add-user',
    help="Add a new user to the system.")
add_user_parser.add_argument('username', help="The username of the new user.")
add_user_parser.add_argument('email', help="The email address of the new user.")
add_user_parser.add_argument('-z', '--timezone', required=False, default=None,
    help="The timezone the user is in. Defaults to UTC.")
add_user_parser.set_defaults(func=add_user)

def modify_user(args):
    from app.models import User
    from saq.database import get_db_connection
    from getpass import getpass

    u = User()
    u.username = args.username

    with get_db_connection() as db:
        c = db.cursor()
        c.execute("""SELECT id FROM users WHERE username = %s""", ( u.username, ))
        row = c.fetchone()
        if row is None:
            logging.error("username {0} does not exist".format(u.username))
            sys.exit(1)

        user_id = row[0]

    if args.email is not None:
        u.email = args.email

    if args.password:
        password = getpass("Enter password for {0}: ".format(u.username))
        confirm = getpass("Confirm password for {0}: ".format(u.username))
        if password != confirm:
            logging.error("passwords do not match")
            sys.exit(1)

        u.password = password

    try:
        import pytz
        u.timezone = 'Etc/UTC'
        if args.timezone:
            u.timezone = args.timezone

        pytz.timezone(u.timezone)

    except Exception as e:
        print(f"ERROR: invalid timezone {target_timezone}")
        sys.exit(1)

    with get_db_connection() as db:
        c = db.cursor()
        if args.email is not None:
            c.execute("""UPDATE users SET email = %s WHERE id = %s""", ( u.email, user_id ))
            
        if args.password:
            c.execute("""UPDATE users SET password_hash = %s WHERE id = %s""", ( u.password_hash, user_id ))

        if args.timezone:
            c.execute("""UPDATE users SET timezone = %s WHERE id = %s""", ( u.timezone, user_id ))

        db.commit()

    logging.info("modified user {0}".format(u.username))

modify_user_parser = subparsers.add_parser('modify-user',
    help="Modifies an existing user on the system.")
modify_user_parser.add_argument('username', help="The username of the user to modify.")
modify_user_parser.add_argument('-e', '--email', dest='email', default=None, help="The new email address of the user.")
modify_user_parser.add_argument('-p', '--password', action='store_true', dest='password', default=False, help="Prompt for a new password.")
modify_user_parser.add_argument('-z', '--timezone', required=False, default=None, help="The timezone the user is in. Defaults to UTC.")
modify_user_parser.set_defaults(func=modify_user)

def delete_user(args):
    from saq.database import get_db_connection

    with get_db_connection() as db:
        c = db.cursor()
        c.execute("""DELETE FROM users WHERE username = %s""", ( args.username, ))
        db.commit()

    logging.info("deleted user {0}".format(args.username))

delete_user_parser = subparsers.add_parser('delete-user',
    help="Deletes an existing user from the system.")
delete_user_parser.add_argument('username', help="The username of the user to modify.")
delete_user_parser.set_defaults(func=delete_user)

# ============================================================================
# test utilities
#

def create_alert(args):
    from saq.analysis import RootAnalysis

    root = RootAnalysis()
    root.tool = 'command line'
    root.tool_instance = 'n/a'
    root.root_type = 'debug'
    root.description = 'Manual Alert'
    root.event_time = datetime.datetime.now()
    root.storage_dir = args.dir

    root.initialize_storage()

    root.details = { 'description': 'manually created root' }
    root.save()

create_alert_parser = subparsers.add_parser('create-alert',
    help="Create a blank alert in the given directory.")
create_alert_parser.add_argument('dir', help="The directory to store the alert in.")
create_alert_parser.set_defaults(func=create_alert)

def test_condition_reporting(args):
    from saq.error import report_condition

    report_condition("SOMETHING ON FIIIRAAAHHH", "j/k")
    time.sleep(5)
    report_condition("SOMETHING ON FIIIRAAAHHH", "j/k")

test_condition_reporting_parser = subparsers.add_parser('test-condition-reporting',
    help="Test the condition reporting system.")
#test_condition_reporting_parser.add_argument('dir', help="The directory to store the alert in.")
test_condition_reporting_parser.set_defaults(func=test_condition_reporting)

def test_error_reporting(args):
    from saq.error import report_exception
    try:
        1 / 0
    except Exception as e:
        report_exception()

test_error_reporting_parser = subparsers.add_parser('test-error-reporting',
    help="Throw an exception and report it.")
test_error_reporting_parser.set_defaults(func=test_error_reporting)

def test_database(args):
    from saq.database import DatabaseSession, User, get_db_connection
    session = DatabaseSession()

    me = session.query(User).filter_by(username = 'jdavison').one()
    logging.debug(str(me))

    with get_db_connection() as db:
        c = db.cursor()
        c.execute("SELECT 1")

    try:
        with get_db_connection() as db:
            raise ValueError("canary")
    except ValueError as e:
        pass

    sys.exit(0)

test_database_parser = subparsers.add_parser('test-database',
    help="Run tests on the database.")
test_database_parser.set_defaults(func=test_database)

def test_database_connections(args):
    import saq
    from saq.database import get_db_connection

    for key in saq.CONFIG.keys():
        if key.startswith('database_'):
            if not saq.CONFIG[key]['hostname'] \
            or not saq.CONFIG[key]['database'] \
            or not saq.CONFIG[key]['username'] \
            or not saq.CONFIG[key]['password']:
                print("skipping {}".format(key))
                continue

            db_name = key[len('database_'):]
            print("trying {}...".format(db_name), end='', flush=True)
            try:
                with get_db_connection(db_name) as db:
                    c = db.cursor()
                    c.execute("SELECT 1")
                    row = c.fetchone()
                    print("OK")
            except Exception as e:
                print("FAILED: {}".format(e))

    sys.exit(0)

test_database_connections_parser = subparsers.add_parser('test-database-connections',
    help="Test the connections to all configured databases.")
test_database_connections_parser.set_defaults(func=test_database_connections)

def test_network_semaphore(args):
    from saq.network_semaphore import NetworkSemaphoreClient
    import time

    client = NetworkSemaphoreClient()
    if client.acquire(args.semaphore_name):
        time.sleep(args.timeout)
        client.release()
    else:
        logging.error("test failed")

network_semaphore_test = subparsers.add_parser('test-network-semaphore',
    help="Test the Network Semaphore Server by requesting a semaphore.")
network_semaphore_test.add_argument('semaphore_name', help="The name of the semaphore to acquire.")
network_semaphore_test.add_argument('-t', '--timeout', required=False, default=60, type=int, dest='timeout',
    help="The number of seconds to wait until the semaphore is released.  Defaults to 60.")
network_semaphore_test.set_defaults(func=test_network_semaphore)

# ============================================================================
# alert management
#

def rebuild_index(args):
    """Rebuilds the indexes for the given alerts."""
    from saq.database import Alert, get_db_connection

    storage_dirs = []
    if args.resync_all:
        with get_db_connection() as db:
            c = db.cursor()
            c.execute("""SELECT storage_dir FROM alerts WHERE location = %s""", (saq.SAQ_NODE,))
            for row in c:
                storage_dirs.append(row[0])
    else:
        storage_dirs = args.dirs

    logging.info("rebuilding indexes for {} alerts".format(len(storage_dirs)))

    for storage_dir in storage_dirs:
        logging.info("rebuilding {}".format(storage_dir))
        alert = saq.db.query(Alert).filter(Alert.storage_dir==storage_dir).first()
        if alert is None:
            logging.error(f"missing alert with storage directory {storage_dir}")
            continue

        try:
            if not alert.load():
                logging.error("unable to load {}".format(alert))
                continue

            alert.rebuild_index()

        except Exception as e:
            logging.error("rebuild failure on {}: {} ({})".format(storage_dir, e, type(e)))
            continue

        finally:
            saq.db.commit()

    sys.exit(0)

rebuild_index_parser = subparsers.add_parser('rebuild-index',
    help="Rebuilds the indexes for the given alerts.")
rebuild_index_parser.add_argument('--all', default=False, action='store_true', dest='resync_all',
    help="Resyncs all alerts that belong to this node. This can take a long time.")
rebuild_index_parser.add_argument('dirs', nargs='*', default=[], help="One ore more alert directories to resync.")
rebuild_index_parser.set_defaults(func=rebuild_index)

def import_alerts(args):
    """Imports one or more alerts from the given directories."""
    import saq
    from saq.constants import ANALYSIS_MODE_CORRELATION
    from saq.database import Alert, DatabaseSession

    for _dir in args.dirs:
        json_path = os.path.join(_dir, 'data.json')
        if not os.path.exists(json_path):
            logging.error("{} does not exist".format(json_path))
            continue

        # load the alert
        alert = Alert()
        alert.storage_dir = _dir
        if not alert.load():
            logging.error("unable to load {}: try running saq upgrade {}".format(_dir, _dir))
            continue

        # has this already already been imported?
        dest_dir = os.path.join(saq.SAQ_HOME, saq.CONFIG['global']['data_dir'], saq.SAQ_NODE, alert.uuid[0:3], alert.uuid)
        if os.path.exists(dest_dir):
            logging.error("ACE storage directory {} already exists".format(dest_dir))
            continue

        try:
            # copy that directory over
            shutil.copytree(_dir, dest_dir)
        except Exception as e:
            logging.error("unable to copy {0} to {1}: {2}".format(
                _dir, dest_dir, str(e)))
            continue

        # remove the old database id if it has one
        alert.id = None
        # and change the storage area
        alert.storage_dir = os.path.relpath(dest_dir, start=saq.SAQ_HOME)

        # are we resetting the alerts?
        if args.reset:
            alert.reset()

        # change a few more things
        alert.location = saq.SAQ_NODE
        alert.company_id = saq.CONFIG['global'].getint('company_id')
        alert.company_name = saq.CONFIG['global']['company_name']
        alert.analysis_mode = ANALYSIS_MODE_CORRELATION
        alert.disposition = None
        alert.disposition_user_id = None
        alert.disposition_time = None
        alert.owner_id = None
        alert.owner_time = None

        # sync it to the database
        alert.sync()

        # request analysis
        alert.schedule()

        logging.info("imported alert {}".format(alert))

import_alert_parser = subparsers.add_parser('import-alerts',
    help="Import one or more alert directories.")
import_alert_parser.add_argument('-r', '--reset', action='store_true', default=False, dest='reset',
    help="Reset imported alerts.")
import_alert_parser.add_argument('dirs', nargs='+', default=[], help="One ore more alert directories to import.")
import_alert_parser.set_defaults(func=import_alerts)

# XXX fix me
def forward_alerts(args):
    """Imports one or more alerts from the given directories."""
    import saq
    from saq.network_client import submit_alerts

    # just make sure each directory has the data.json file in it
    for _dir in args.dirs:
        json_path = os.path.join(_dir, 'data.json')
        if not os.path.exists(json_path):
            logging.error("{} does not exist".format(json_path))
            continue

    try:
        # submit them all at once
        submit_alerts(args.remote_host, args.remote_port,
                      args.ssl_cert, args.ssl_hostname, args.ssl_key, args.ca_path, args.dirs)
    except Exception as e:
        logging.error("alert submission failed: {}".format(e))
        traceback.print_exc()
        sys.exit(1)

    logging.info("submitted {} alerts".format(len(args.dirs)))
    sys.exit(0)

forward_alert_parser = subparsers.add_parser('forward-alerts',
    help="Forward one or more alerts to another ACE system.")
forward_alert_parser.add_argument('-r', '--remote-host', dest='remote_host', required=True,
    help="Network location of remote ACE system.")
forward_alert_parser.add_argument('-p', '--remote-port', dest='remote_port', type=int, required=True,
    help="Network port of remote ACE system.")
forward_alert_parser.add_argument('--ssl-cert', dest='ssl_cert', required=True,
    help="Path to the SSL certificate to use when submitting (SSL parameter.)")
forward_alert_parser.add_argument('--ssl-key', dest='ssl_key', required=True,
    help="Path to the SSL key to use when submitting (SSL parameter.)")
forward_alert_parser.add_argument('--ssl-hostname', dest='ssl_hostname', required=True,
    help="The common name expected from the SSL server (SSL parameter).  This may be different than the hostname.")
forward_alert_parser.add_argument('--ca-path', dest='ca_path', required=True,
    help="Path to the root certificate chain to use when submitting (SSL parameter.)")
forward_alert_parser.add_argument('dirs', nargs='+', default=[], help="One ore more alert directories to forward.")
forward_alert_parser.set_defaults(func=forward_alerts)

def delete_alerts(args):
    """Completely deletes the given alerts from both the storage system and the database."""
    import saq
    from saq.database import Alert, DatabaseSession

    for uuid in args.uuids:
        try:
            # we do them one at a time in case one of them fails
            session = DatabaseSession()
            session.execute(Alert.__table__.delete().where(Alert.uuid == uuid))
            session.commit()
            session.close()
        except Exception as e:
            logging.error("unable to delete alert {0}: {1}".format(uuid, str(e)))

    for uuid in args.uuids:
        storage_dir = os.path.join(saq.SAQ_HOME, saq.CONFIG['global']['data_dir'], saq.CONFIG['global']['node'], uuid[0:3], uuid)
        if not os.path.exists(storage_dir):
            logging.warning("storage directory {0} does not exist".format(storage_dir))
            continue

        try:
            shutil.rmtree(storage_dir)
        except Exception as e:
            logging.error("unable to delete storage directory {0}: {1}".format(storage_dir, str(e)))

    sys.exit(0)

delete_alert_parser = subparsers.add_parser('delete-alerts',
    help="Delete one or more alerts by UUID.")
delete_alert_parser.add_argument('uuids', nargs='+', default=[], help="One ore more alert UUIDs to delete.")
delete_alert_parser.set_defaults(func=delete_alerts)

def reset_alerts(args): 
    import saq
    from saq.analysis import RootAnalysis
    from saq.database import Alert, DatabaseSession

    for storage_dir in args.dirs:
        # get the storage directory of the alert
        if not os.path.exists(storage_dir):
            logging.error("storage directory {0} does not exist".format(storage_dir))
            continue

        session = None

        # try to load it from the database first
        try:
            session = DatabaseSession()
            root = session.query(Alert).filter(Alert.storage_dir==storage_dir).one()
            logging.info("loaded {} from database".format(storage_dir))
        except:
            root = RootAnalysis()
            root.storage_dir = storage_dir
        finally:
            if session:
                session.close()

        try:
            root.load()
        except Exception as e:
            logging.error("unable to load {}: {}".format(root.storage_dir, e))
            continue

        root.reset()
        root.save()

# reset-alerts
reset_alert_parser = subparsers.add_parser('reset-alerts',
    help="Reset the given alerts allowing for re-analysis.")
reset_alert_parser.add_argument('dirs', nargs='+', help="One or more alert directories to reset.")
reset_alert_parser.set_defaults(func=reset_alerts)

def archive_alerts(args):
    import saq
    from saq.analysis import RootAnalysis
    from saq.database import Alert, DatabaseSession

    for storage_dir in args.dirs:
        # get the storage directory of the alert
        if not os.path.exists(storage_dir):
            logging.error("storage directory {} does not exist".format(storage_dir))
            continue

        alert = saq.db.session.query(Alert).filter(Alert.storage_dir==storage_dir).first()
        if alert is None:
            logging.warning(f"cannot find alert with storage_dir {storage_dir}")
            continue

        try:
            alert.load()
        except Exception as e:
            logging.error("unable to load {}: {}".format(root.storage_dir, e))
            continue

        alert.archive()
        alert.save()

        # mark the alert as archived
        #try:
            #saq.db.session.execute(Alert.__table__.update().where(Alert.storage_dir==storage_dir).values(
                #Alert.archived = True))
            #saq.db.session.commit()

            #session = DatabaseSession()
            #root = saq.db.session.query(Alert).filter(Alert.storage_dir==storage_dir).one()
            #root.archived = True
            #saq.db.session.add(root)
            #saq.db.session.commit()

        #except Exception as e:
            #logging.error("unable to mark alert {} as archived: {}".format(root, e))
        #finally:
            #if session:
                #session.close()

# reset-alerts
reset_alert_parser = subparsers.add_parser('archive-alerts',
    help="Archives a given alert by deleting analysis details and external files but keeping observations and tags.")
reset_alert_parser.add_argument('dirs', nargs='+', help="One or more alert directories to archive.")
reset_alert_parser.set_defaults(func=archive_alerts)

def add_observable(args):
    import saq
    import saq.constants
    from saq.analysis import RootAnalysis
    from saq.lock import initialize_locking

    initialize_locking()

    if args.observable_type not in saq.constants.VALID_OBSERVABLE_TYPES:
        logging.error("invalid observable type {}".format(args.observable_type))
        sys.exit(1)

    # get the alert to modify
    alert = RootAnalysis()
    alert.storage_dir = args.dir
    if not alert.lock():
        logging.error("unable to lock alert {}".format(alert))
        sys.exit(1)

    try:
        alert.load()

        if args.observable_type == saq.constants.F_FILE or args.observable_type == saq.constants.F_SUSPECT_FILE:
            try:
                dest_path = os.path.join(alert.storage_dir, os.path.basename(args.observable_value))
                shutil.copy(args.observable_value, dest_path)
                args.observable_value = os.path.relpath(dest_path, start=alert.storage_dir)
            except Exception as e:
                logging.error("unable to copy file into storage directory: {0}".format(str(e)))

        alert.add_observable(args.observable_type, args.observable_value, o_time=args.reference_time)
        alert.save()

    except Exception as e:
        logging.error(str(e))
        traceback.print_exc()
        sys.exit(1)

    finally:
        alert.unlock()

# add-observable
add_observable_parser = subparsers.add_parser('add-observable',
    help="Add an observable to an existing alert and re-analyze.")
add_observable_parser.add_argument('dir', help="The path to the alert to modify.")
add_observable_parser.add_argument('observable_type', help="The type of the observable to add.")
add_observable_parser.add_argument('observable_value', help="The value of the observable.")
add_observable_parser.add_argument('-t', '--reference-time', required=False, dest='reference_time', default=None,
    help="Specify a datetime in YYYY-MM-DD HH:MM:SS format the observable should be referenced from.")
add_observable_parser.set_defaults(func=add_observable)

def reload_alerts(args):
    from saq.constants import ANALYSIS_MODE_CORRELATION
    from saq.database import Alert, DatabaseSession

    # generate the list of alerts to reload
    session = DatabaseSession()
    for uuid in args.uuids:
        alert = session.query(Alert).filter(Alert.uuid == uuid).one()
        #alert.request_correlation()
        alert.analysis_mode = ANALYSIS_MODE_CORRELATION
        alert.schedule()

reload_alert_parser = subparsers.add_parser('reload-alerts',
    help="Force analysis (again) on one or more existing alert(s).")
reload_alert_parser.add_argument('uuids', nargs='+',
    help="One or more alert UUIDs to analyze.")
reload_alert_parser.set_defaults(func=reload_alerts)

# ============================================================================
# general utility commands
#

#
# messaging
#

def display_messages(args):
    from saq.database import Message, MessageRouting
    from sqlalchemy.orm import joinedload
    for message in saq.db.query(Message).options(joinedload('routing')).order_by(Message.id):
        print(message.content)
        for message_route in message.routing:
            print(f"\t{message_route.route}\t{message_route.destination}")

    sys.exit(0)

display_messages_parser = subparsers.add_parser('display-messages',
    help="Displays any queued messages that have not been sent yet.")
display_messages_parser.set_defaults(func=display_messages)

def send_message(args):
    from saq.messaging import send_message, initialize_message_system
    initialize_message_system()
    send_message(args.message, args.message_type)
    sys.exit(0)

send_message_parser = subparsers.add_parser('send-message',
    help="Sends the given message (and type) to the messaging system.")
send_message_parser.add_argument('message',
    help="The plain text message to send.")
send_message_parser.add_argument('-t', '--message-type', required=False, default=None,
    help="The type of the message (determines the route.) Default sends to all routes.")
send_message_parser.set_defaults(func=send_message)

def dispatch_messages(args):
    from saq.messaging import initialize_message_system, start_message_system, wait_message_system
    initialize_message_system()
    start_message_system()
    wait_message_system()
    sys.exit(0)

dispatch_messages_parser = subparsers.add_parser('dispatch-messages',
    help="Dispatches all the messages in the message queue.")
dispatch_messages_parser.set_defaults(func=dispatch_messages)

def clear_messages(args):
    from saq.database import Message, MessageRouting
    # if we specified a route then we clear all the messages sent to that route that match the given content
    if args.route is not None:
        sql = MessageRouting.__table__.delete().where(MessageRouting.route == args.route)
        if args.message_content is not None:
            sql = sql.where(MessageRouting.message_id.in_(saq.db.query(Message.id).filter(Message.content.like('%{}%'.format(args.message_content)))))
    else:
        if args.message_content is not None:
            sql = Message.__table__.delete().where(Message.content.like('%{}%'.format(args.message_content)))
        else:
            sql = Message.__table__.delete()

    saq.db.execute(sql)
    saq.db.execute(Message.__table__.delete().where(Message.id.notin_(saq.db.query(MessageRouting.message_id))))
    saq.db.commit()
    sys.exit(0)

clear_messages_parser = subparsers.add_parser('clear-messages',
    help="Clears all or part of the message queue, depending on the options. By default all messages are cleared out.")
clear_messages_parser.add_argument('-r', '--route', required=False, default=None,
    help="Limit to messages sent to the given route.")
clear_messages_parser.add_argument('-m', '--message-content', required=False, default=None,
    help="Limit to messages that have content that matches the this option.")
clear_messages_parser.set_defaults(func=clear_messages)

def submit_failed_submissions(args):
    import saq
    from saq.network_client import submit_alerts

    failed_dir = os.path.join(saq.SAQ_HOME, saq.CONFIG['network_client_ace']['failed_dir'])
    if not os.path.isdir(failed_dir):
        logging.warning("{} does not exist (no submissions available?".format(failed_dir))
        sys.exit(0)

    storage_dirs = [os.path.join(failed_dir, d) for d in os.listdir(failed_dir)]
    if not storage_dirs:
        logging.warning("nothing available to submit")
        sys.exit(0)

    logging.info("{} submissions available".format(len(storage_dirs)))

    try:
        remote_host = saq.CONFIG['network_client_ace']['remote_host']
        remote_port = saq.CONFIG['network_client_ace'].getint('remote_port')
        ssl_cert = os.path.join(saq.SAQ_HOME, saq.CONFIG['network_client_ace']['ssl_cert'])
        ssl_key = os.path.join(saq.SAQ_HOME, saq.CONFIG['network_client_ace']['ssl_key'])
        ca_path = os.path.join(saq.SAQ_HOME, saq.CONFIG['network_client_ace']['ca_path'])
        remote_hostname = saq.CONFIG['network_client_ace']['ssl_hostname']

        submit_alerts(remote_host, remote_port, ssl_cert, remote_hostname, ssl_key, ca_path, storage_dirs)

        if args.remove:
            for storage_dir in storage_dirs:
                try:
                    shutil.rmtree(storage_dir)
                    logging.info("removed {}".format(storage_dir))
                except Exception as e:
                    logging.error("unable to remove {}: {}".format(storage_dir))

        sys.exit(0)

    except Exception as e:
        logging.error("unable to submit: {}".format(e))
        sys.exit(1)

submit_failed_submissions_parser = subparsers.add_parser('submit-failed-submissions',
    help="Submit any failed submissions stored in the directory defined in the network_client_ace configuration section.")
submit_failed_submissions_parser.add_argument('-r', '--remove', required=False, default=False, action='store_true', dest='remove',
    help="Remove the directories that were submitted after a succuessful submissions.")
submit_failed_submissions_parser.set_defaults(func=submit_failed_submissions)

def search_archive(args):
    import saq
    from saq.database import _get_db_connection

    # are we exporting into a directory?
    if args.output_dir:
        if not os.path.isdir(args.output_dir):
            try:
                os.mkdir(args.output_dir)
            except Exception as e:
                logging.error("unable to create output directory {}: {}".format(args.output_dir, e))
                sys.exit(0)

    search_items = args.search_items
    
    # are we reading search from standard input?
    if args.from_stdin:
        for search_item in sys.stdin:
            search_items.append(search_item.strip())

    db = _get_db_connection('email_archive')
    c = db.cursor()
    query = """
SELECT
    archive_server.hostname, HEX(archive.md5)
FROM
    archive JOIN archive_server ON archive.server_id = archive_server.server_id
    {extended_from_clause}
WHERE 
    {where_clauses}
"""
    
    where_clauses = []
    parameters = []

    extended_from_clause = "JOIN archive_search ON archive.archive_id = archive_search.archive_id"
    if args.exact:
        extended_from_clause = "JOIN archive_index ON archive.archive_id = archive_index.archive_id"
    
    for search_item in search_items:
        if not any([args.env_from, args.env_to, args.mail_from, args.mail_to, args.subject, args.message_id]):
            if args.exact:
                where_clauses.append("archive_index.hash = UNHEX(MD5(%s))")
                parameters.append(search_item)
            else:
                where_clauses.append("archive_search.value LIKE %s")
                parameters.append('%%{}%%'.format(search_item))
        else:
            _archive_index_template = "(archive_index.field = '{field}' AND archive_index.hash = UNHEX(MD5(%s)))"
            _archive_index_value = search_item
            _archive_search_template = "(archive_search.field = '{field}' AND archive_search.value LIKE %s)"
            _archive_search_value = '%%{}%%'.format(search_item)

            if args.env_from:
                if args.exact:
                    where_clauses.append(_archive_index_template.format(field='env_from'))
                    parameters.append(_archive_index_value)
                else:
                    where_clauses.append(_archive_search_template.format(field='env_from'))
                    parameters.append(_archive_search_value)

            if args.env_to:
                if args.exact:
                    where_clauses.append(_archive_index_template.format(field='env_to'))
                    parameters.append(_archive_index_value)
                else:
                    where_clauses.append(_archive_search_template.format(field='env_to'))
                    parameters.append(_archive_search_value)

            if args.mail_from:
                if args.exact:
                    where_clauses.append(_archive_index_template.format(field='mail_from'))
                    parameters.append(_archive_index_value)
                else:
                    where_clauses.append(_archive_search_template.format(field='mail_from'))
                    parameters.append(_archive_search_value)

            if args.mail_to:
                if args.exact:
                    where_clauses.append(_archive_index_template.format(field='mail_to'))
                    parameters.append(_archive_index_value)
                else:
                    where_clauses.append(_archive_search_template.format(field='mail_to'))
                    parameters.append(_archive_search_value)

            if args.subject:
                if args.exact:
                    where_clauses.append(_archive_index_template.format(field='subject'))
                    parameters.append(_archive_index_value)
                else:
                    where_clauses.append(_archive_search_template.format(field='subject'))
                    parameters.append(_archive_search_value)

            if args.url:
                if args.exact:
                    where_clauses.append(_archive_index_template.format(field='url'))
                    parameters.append(_archive_index_value)
                else:
                    where_clauses.append(_archive_search_template.format(field='url'))
                    parameters.append(_archive_search_value)

            if args.message_id:
                if args.exact:
                    where_clauses.append(_archive_index_template.format(field='message_id'))
                    parameters.append(_archive_index_value)
                else:
                    where_clauses.append(_archive_search_template.format(field='message_id'))
                    parameters.append(_archive_search_value)

    query = query.format(extended_from_clause=extended_from_clause, 
                         where_clauses=' OR '.join(where_clauses))
    #print(query)
    #print(','.join(parameters))
    c.execute(query, parameters)
    for server, md5 in c:
        # does this archive file exist?
        archive_base_dir = os.path.join(saq.DATA_DIR, saq.CONFIG['analysis_module_email_archiver']['archive_dir'])
        if args.archive_dir:
            archive_base_dir = args.archive_dir

        if not os.path.isdir(archive_base_dir):
            logging.error("archive directory {} does not exist".format(archive_base_dir))
            sys.exit(1)

        archive_path = os.path.join(archive_base_dir, server.lower(), md5.lower()[0:3], 
                                    '{}.gz.gpg'.format(md5.lower()))
        if not os.path.exists(archive_path):
            logging.warning("archive email {} does not exist at {}".format(md5, archive_path))
            continue

        # are we just listing the paths?
        if not args.output_dir:
            print(archive_path)
            continue

        # are we just copying the entire encrypted email?
        if not saq.ENCRYPTION_PASSWORD:
            try:
                shutil.copy(archive_path, args.output_dir)
            except Exception as e:
                logging.warning("unable to copy {} to {}: {}".format(archive_path, args.output_dir, e))
                continue

            print(os.path.join(args.output_dir, os.path.basename(archive_path)))
            continue

        # decrypt and decompress the emails
        dest_path = os.path.join(args.output_dir, md5.lower())
        with open(dest_path, 'wb') as fp:
            gpg_p = Popen(['gpg', '--no-tty', '-d', '--passphrase-fd', '0', archive_path], 
                          stdout=PIPE, stderr=PIPE, stdin=PIPE)
            gunzip_p = Popen(['zcat'], stdin=gpg_p.stdout, stdout=fp)
            gpg_p.stdin.write('{}\n'.format(saq.ENCRYPTION_PASSWORD).encode())
            gpg_p.stdin.close()
            gpg_p.stdout.close()
            gunzip_p.communicate()
            gpg_p.wait()

            if gpg_p.returncode != 0:
                logging.warning("gpg decryption failed for {} (return code {})".format(archive_path, gpg_p.returncode))
                continue

        print(dest_path)

    db.close()
    sys.exit(0)

# search-archive
search_archive_parser = subparsers.add_parser('search-archive',
    help="Search the email archives.")
search_archive_parser.add_argument('--env-from', required=False, default=False, action='store_true', dest='env_from',
    help="Narrow searching to the envelope MAIL FROM field.")
search_archive_parser.add_argument('--env-to', required=False, default=False, action='store_true', dest='env_to',
    help="Narrow searching to the envelope RCPT TO field.")
search_archive_parser.add_argument('--mail-from', required=False, default=False, action='store_true', dest='mail_from',
    help="Narrow searching to the message body From field.")
search_archive_parser.add_argument('--mail-to', required=False, default=False, action='store_true', dest='mail_to',
    help="Narrow searching to the message body To field.")
search_archive_parser.add_argument('--subject', required=False, default=False, action='store_true', dest='subject',
    help="Narrow searching to the message body Subject field.")
search_archive_parser.add_argument('--message-id', required=False, default=False, action='store_true', dest='message_id',
    help="Narrow searching to the message body Message-ID field.")
search_archive_parser.add_argument('--url', required=False, default=False, action='store_true', dest='url',
    help="Narrow searching to a URL found anywhere in the email.")
search_archive_parser.add_argument('--exact', required=False, default=False, action='store_true', dest='exact',
    help="Perform an exact match (ignores options to narrow down search.) This is the fastest search.")
search_archive_parser.add_argument('--from-stdin', required=False, default=False, action='store_true', dest='from_stdin',
    help="Read search items from standard input (one per line.)")
search_archive_parser.add_argument('-d', '--output-dir', required=False, default=None, dest='output_dir',
    help="Export selected emails into the given directory. Also see the -p option in the main command.")
search_archive_parser.add_argument('-a', '--archive-dir', required=False, default=None, dest='archive_dir',
    help="Specify an alternative email archive directory. Defaults to what is specified in the analysis_module_email_archiver configuration.")
search_archive_parser.add_argument('search_items', nargs='*',
    help="One or more things to search for.  Each query will be searhed for individually.")
search_archive_parser.set_defaults(func=search_archive)

#
# remediation
#

def display_remediation_requests(args):
    import saq
    from saq.remediation.constants import REMEDIATION_STATUS_NEW, REMEDIATION_STATUS_IN_PROGRESS
    from saq.database import Remediation

    #print(['ID', 'STATUS', 'TYPE', 'ACTION', 'DATE', 'TARGET']))
    row_format = "{:<8}{:<10}{:<8}{:<9}{:<20} {}"
    print(row_format.format('ID', 'STATUS', 'TYPE', 'ACTION', 'DATE', 'TARGET'))
    for r in saq.db.query(Remediation).filter(Remediation.company_id == saq.COMPANY_ID,
                                              Remediation.status.in_([REMEDIATION_STATUS_NEW, 
                                                                      REMEDIATION_STATUS_IN_PROGRESS]))\
                                      .order_by(Remediation.id):

        print(row_format.format(r.id, r.status, r.type, r.action, str(r.insert_date), r.key))
        #print('\t'.join(map(str, [r.id, r.status, r.type, r.action, r.insert_date, r.key])))

    sys.exit(0)

display_remediation_parser = subparsers.add_parser('display-remediation-requests', 
    help="Displays the remediation requests currently in the queue or in processing.")
display_remediation_parser.set_defaults(func=display_remediation_requests)

def clear_remediation_request(args):
    import saq
    from saq.database import Remediation
    from saq.remediation import REMEDIATION_STATUS_NEW
    from sqlalchemy import and_

    if not args.remediation_ids and not args.all:
        logging.error("no remediation ids were specified and the --all option was not used")
        sys.exit(1)
    
    clause = and_(Remediation.status == REMEDIATION_STATUS_NEW,
                  Remediation.lock == None,
                  Remediation.company_id == saq.COMPANY_ID)
    if not args.all and args.remediation_ids:
        clause = and_(clause, Remediation.id.in_(args.remediation_ids))

    logging.info("deleted {} requests".format(
        saq.db.execute(Remediation.__table__.delete().where(clause)).rowcount))
    
    saq.db.commit()
    sys.exit(0)

clear_remediation_request_parser = subparsers.add_parser('clear-remediation-requests',
    help="Clears one or more remediation requests.")
clear_remediation_request_parser.add_argument('-a', '--all', action='store_true', default=False,
    help="Clears all remediation requests that are not locked or have expired locks.")
clear_remediation_request_parser.add_argument('remediation_ids', nargs=argparse.REMAINDER,
    help="Zero or more remediation IDs to clear which can be obtained using the display-remediation-request command.")
clear_remediation_request_parser.set_defaults(func=clear_remediation_request)

def remediate_emails(args):
    from saq.database import User
    from saq.remediation import get_remediation_targets, initialize_remediation_system_manager, request_remediation, \
                                REMEDIATION_TYPE_EMAIL, request_restoration

    initialize_remediation_system_manager()

    user_id = saq.db.query(User).filter(User.username == args.user_name).one().id

    if args.from_stdin:
        args.targets.extend([_.strip() for _ in sys.stdin])

    if args.message_id_only:
        remediation_targets = get_remediation_targets(args.targets)
    else:
        targets = args.targets[:]
        remediation_targets = []
        while targets:
            recipient = targets.pop()
            message_id = targets.pop()
            target_func = request_remediation
            if args.restore:
                target_func = request_restoration
            result = target_func(REMEDIATION_TYPE_EMAIL, message_id, recipient, 
                                         user_id=user_id, comment=args.comment, company_id=saq.COMPANY_ID)
            print(f"got result {result}")

    sys.exit(0)

remediate_email_parser = subparsers.add_parser('remediate-emails',
    help="Remediate the given emails by message-id and recipient.")
remediate_email_parser.add_argument('--restore', required=False, action='store_true', default=False,
    help="Restore the given emails instead of removing them.")
remediate_email_parser.add_argument('--from-stdin', required=False, dest='from_stdin', action='store_true', default=False,
    help="Read the message-ids and/or recipients from standard input.")
remediate_email_parser.add_argument('-m', '--message-id-only', required=False, default=False, action='store_true',
    help="Assume all parameters are message-ids. Use email archive database to determine recipients automatically.")
remediate_email_parser.add_argument('--remediation-system', required=False, default='phishfry', choices=['phishfry'],
    help="Select which remediation system to use. Defaults to phishfry.")
remediate_email_parser.add_argument('-c', '--comment', required=False, default=None,
    help="An optional comment to add to the remediation.")
remediate_email_parser.add_argument('-u', '--user-name', required=False,
    help="The username to execute the remediation as. Defaults to the current user name.")
remediate_email_parser.add_argument('targets', nargs='*',
    help="One or more message-ids to remediate. You can also specify --from-stdin.")
remediate_email_parser.set_defaults(func=remediate_emails)

def remediate_email_old(args):
    from saq.remediation import _create_remediation_email, _remediate_email_lotus, _remediate_email_o365
    message_ids = args.message_ids
    if args.from_stdin:
        for message_id in sys.stdin:
            message_ids.append(message_id.strip())

    logging.debug("remediating {} emails")
    
    env_mail_from = args.env_mail_from
    env_rcpt_to = []
    if args.env_rcpt_to:
        env_rcpt_to.append(args.env_rcpt_to)
    mail_from = args.mail_from
    mail_to = []
    if args.mail_to:
        mail_to.append(args.mail_to)
    subject = args.subject
    decoded_subject = args.decoded_subject

    for message_id in message_ids:
        try:
            _remediate_email_lotus([_create_remediation_email(env_mail_from=env_mail_from,
                                                              env_rcpt_to=env_rcpt_to,
                                                              mail_from=mail_from,
                                                              decoded_mail_from=None,
                                                              mail_to=mail_to,
                                                              subject=subject,
                                                              decoded_subject=decoded_subject,
                                                              message_id=message_id)])
        except Exception as e:
            logging.error("unable to remediate {} in lotus: {}".format(message_id, e))
            report_exception()

    try:
        _remediate_email_o365(message_ids)
    except Exception as e:
        logging.error("unable to remediate in office365: {}".format(e))
        report_exception()

    sys.exit(0)

# remediate-email
remediate_email_parser = subparsers.add_parser('remediate-email',
    help="Remediate the given emails by message-id.")
remediate_email_parser.add_argument('--from-stdin', required=False, dest='from_stdin', action='store_true', default=False,
    help="Read the message-ids from standard input.")
remediate_email_parser.add_argument('--env-mail-from', required=False, dest='env_mail_from', default=None,
    help="(LOTUS NOTES ONLY) Envelope MAIL FROM")
remediate_email_parser.add_argument('--env-rcpt-to', required=False, dest='env_rcpt_to', default=None,
    help="(LOTUS NOTES ONLY) Envelope RCPT TO")
remediate_email_parser.add_argument('--mail-from', required=False, dest='mail_from', default=None,
    help="(LOTUS NOTES ONLY) Mail From")
remediate_email_parser.add_argument('--mail-to', required=False, dest='mail_to', default=None,
    help="(LOTUS NOTES ONLY) Mail To")
remediate_email_parser.add_argument('--subject', required=False, dest='subject', default=None,
    help="(LOTUS NOTES ONLY) Subject")
remediate_email_parser.add_argument('--decoded-subject', required=False, dest='decoded_subject', default=None,
    help="(LOTUS NOTES ONLY) Subject (Decoded)")
remediate_email_parser.add_argument('message_ids', nargs='*',
    help="One or more message-ids to remediate. You can also specify --from-stdin.")
remediate_email_parser.set_defaults(func=remediate_email_old)

def submit_phishme_response(args):
    from saq.phishme import submit_response
    submit_response(args.recipient, args.subject, args.disposition, args.comment)
    sys.exit(0)

submit_phishme_response_parser = subparsers.add_parser('submit-phishme-response',
    help="Submits a PhishMe response with the given parameters. Useful for debugging purposes.")
submit_phishme_response_parser.add_argument('recipient', help="Target recipient of the response.")
submit_phishme_response_parser.add_argument('subject', help="Original subject of the reported email.")
submit_phishme_response_parser.add_argument('disposition', help="Disposition assigned to the report.")
submit_phishme_response_parser.add_argument('--comment', 
    help="Optional user comment to add to the report.")
submit_phishme_response_parser.set_defaults(func=submit_phishme_response)

def update_organization(args):

    from saq.modules import LDAPAnalysisModule
    import saq

    # load the organization information
    config = saq.CONFIG['analysis_module_user_tagger']

    dest_file = os.path.join(saq.SAQ_HOME, config['json_path'])
    temp_file = os.path.join(saq.SAQ_HOME, '{0}.tmp'.format(config['json_path']))

    # key = userID (lowercase), value = set(tags...)
    mapping = {}

    # horrible copy-pasta (sorry)
    # load ldap settings from configuration file
    ldap_server = saq.CONFIG.get('ldap', 'ldap_server')
    ldap_port = saq.CONFIG.getint('ldap', 'ldap_port') or 389
    ldap_bind_user = saq.CONFIG.get('ldap', 'ldap_bind_user')
    ldap_bind_password = saq.CONFIG.get('ldap', 'ldap_bind_password')
    ldap_base_dn = saq.CONFIG.get('ldap', 'ldap_base_dn')

    def ldap_query(query):

        from ldap3 import Server, Connection, SIMPLE, SYNC, ASYNC, SUBTREE, ALL, ALL_ATTRIBUTES
        import json

        try:
            with Connection(
                Server(ldap_server, port = ldap_port, get_info = ALL), 
                auto_bind = True,
                client_strategy = SYNC,
                user=ldap_bind_user,
                password=ldap_bind_password,
                authentication=SIMPLE, 
                check_names=True) as c:

                logging.debug("running ldap query for ({0})".format(query))
                c.search(ldap_base_dn, '({0})'.format(query), SUBTREE, attributes = ALL_ATTRIBUTES)

                # a little hack to move the result into json
                response = json.loads(c.response_to_json())
                result = c.result

                if len(response['entries']) < 1:
                    return None

                # XXX not sure about the 0 here, I guess only if we only looking for one thing at a time
                return response['entries'][0]['attributes']

                # look for the result with the 'type' set to 'searchResEntry'
                #for r in response['entries']:
                    #if r['type'] == 'searchResEntry':
                        # and the what we're looking for should be in here
                    #return r['attributes']

                return None

        except Exception as e:
            logging.error("unable to perform ldap query: {0}".format(str(e)))
            report_exception()
            return None

    if os.path.exists(temp_file):
        try:
            os.remove(temp_file)
        except Exception as e:
            logging.error("unable to remove temp file {0}".format(temp_file))
            sys.exit(1)

    def recurse_org(group_name, limit, tag, current_user_id, current_level=0):
        # add this user
        if current_user_id.lower() not in mapping:
            mapping[current_user_id.lower()] = set()

        mapping[current_user_id.lower()].add(tag)

        current_level += 1
        if limit != 'all' and current_level > int(limit):
            return

        # figure out who works for this guy
        query_results = ldap_query("cn={0}*".format(current_user_id))
        if query_results is None:
            return

        if 'directReports' in query_results:
            for direct_report in query_results['directReports']:
                # extract the userID from this thing
                m = re.search(r'CN=([^,]+),', direct_report)
                if m is None:
                    logging.warning("unable to extract direct report info from {0}".format(direct_report))

                user_id = m.group(1)
                if user_id is not None:
                    # add this guy (and possibly all his direct reports too)
                    logging.debug("adding {0} as a direct report to {1} for group {2}".format(user_id, current_user_id, group_name))
                    recurse_org(group_name, limit, tag, user_id, current_level)

    # load all the hierarchy definitions
    for section in config.keys():
        if section.startswith('group_'):
            m = re.match(r'^group_([^_]+)$', section)
            if m is None:
                logging.error("unable to parse group name from {0}".format(section))
                continue

            group_name = m.group(1)
            parent_id, limit, tag = [x.strip() for x in config[section].split(',')]
            
            recurse_org(group_name, limit, tag, parent_id)

    # write out the json
    for key in mapping.keys():
        mapping[key] = list(mapping[key])

    with open(temp_file, 'w') as fp:
        json.dump(mapping, fp)

    # finally update the production file
    try:
        shutil.move(temp_file, dest_file)
    except Exception as e:
        logging.error("unable to move {0} to {1}: {2}".fomrat(temp_file, dest_file, str(e)))

update_organization_parsers = subparsers.add_parser('update-organization',
    help="Updates the files used by the UserTaggingAnalyzer module.")
update_organization_parsers.set_defaults(func=update_organization)

def update_crits_cache(args):
    from saq.crits import update_local_cache
    if update_local_cache():
        sys.exit(0)

    sys.exit(1)

update_crits_cache_parser = subparsers.add_parser('update-crits-cache',
    help="Updates the local CRITS cache used by the analysis modules.")
update_crits_cache_parser.set_defaults(func=update_crits_cache)

def update_sip_cache(args):
    from saq.intel import update_local_cache
    if update_local_cache():
        sys.exit(0)

    sys.exit(1)

update_sip_cache_parser = subparsers.add_parser('update-sip-cache',
    help="Updates the local SIP cache used by the analysis modules.")
update_sip_cache_parser.set_defaults(func=update_sip_cache)

def export_sip_yara_rules(args):

    import yara
    yara.set_config(max_strings_per_rule=30720)

    import pysip
    sip_client = pysip.Client(saq.CONFIG['sip']['remote_address'], saq.CONFIG['sip']['api_key'], verify=False)

    from saq.intel import get_indicator_type_mapping 

    # make sure target directory exists
    if not os.path.isdir(args.dir):
        try:
            os.makedirs(args.dir)
        except Exception as e:
            logging.error("unable to create directory {}: {}".format(args.dir, e))
            sys.exit(1)

    # load the mapping from indicator type to the string modifiers to use
    string_modifiers = collections.defaultdict(lambda: saq.CONFIG['sip_yara_export_string_modifiers']['default'])
    for indicator_type in saq.CONFIG['sip_yara_export_string_modifiers']:
        if indicator_type == 'default':
            continue

        string_modifiers[indicator_type] = saq.CONFIG['sip_yara_export_string_modifiers'][indicator_type]
        logging.debug(f"using string modifiers {string_modifiers[indicator_type]} for {indicator_type}")


    # what is the minimum length an indicator value can be to be put into a yara rule?
    export_minimum_length = saq.CONFIG['sip_yara_export'].getint('export_minimum_length')

    def format_yara_string(s):
        return s.replace('\\', '\\\\').replace('"','\\"').replace("\n","")

    def get_yara_filename(indicator_type):
        return os.path.join(args.dir, "SIP_{}.yar".format(indicator_type))

    indicator_types = [_.strip() for _ in saq.CONFIG['sip_yara_export']['export_list'].split(',')]
    included_sources = [_.strip() for _ in saq.CONFIG['sip_yara_export']['export_sources_include'].split(',')]
    excluded_sources = [_.strip() for _ in saq.CONFIG['sip_yara_export']['export_sources_exclude'].split(',')]

    sources = []
    for source in sip_client.get('/api/intel/source'):
        if included_sources and source not in included_sources:
            continue

        if source in excluded_sources:
            continue

        sources.append(source)

    for indicator_type in indicator_types:
        # what is the crits equivalent of this indicator type?
        indicator_type = get_indicator_type_mapping(indicator_type)
        logging.info("exporting indicator type {}".format(indicator_type))
      
        # does a template file exists for this indicator type?
        template_path = os.path.join(saq.CONFIG['sip_yara_export']['export_template_dir'], '{}.template'.format(indicator_type))
        if not os.path.exists(template_path):
            template_path = os.path.join(saq.CONFIG['sip_yara_export']['export_template_dir'], 'default.template')

        logging.debug("using template {} for {}".format(template_path, indicator_type))

        # load the template we're going to be using
        with open(template_path, 'r') as fp:
            template = fp.read()

        template = template.replace('TEMPLATE_RULE_NAME', 'SIP_{}'.format(re.sub(r'[^a-zA-Z0-9_]', '', indicator_type)))
        
        special_paths = { "%temp%":[ "\\windows\\temp", "\\temp", "\\appdata\\local\\temp", "\\local settings\\temp", "\\locals~1\\temp" ],
                          "%appdata%": [ "\\application data", "\\appdata\\roaming" ],
                          "%programdata%": [ "\\programdata", "\\documents and settings\\all users" ],
                          "%programfiles%": [ "\\program files", "\\program files (x86)" ],
                          "%systemdrive%": [ "" ],
                          "%system%": [ "\\windows\\system32", "\\windows\\system" ] }

        string_data = io.StringIO()
        count = 0
        skip_count = 0

        url = f'/api/indicators?status=Analyzed&bulk=True&type={indicator_type}'
        if sources:
            url += '&sources={}'.format(','.join(sources))

        for indicator in sip_client.get(url):
            string_buffer = io.StringIO()
            item_id = indicator['id']
            item_value = indicator['value']

            # is this string too small?
            if len(item_value) < export_minimum_length:
                skip_count += 1
                continue
                
            # do we need to reformat this string?
            if indicator['type'] == get_indicator_type_mapping(saq.intel.I_FILE_PATH):
                subindicator = 0
                for path in special_paths:
                    if path.lower() in indicator['value'].lower():
                        for p_item in special_paths[path]:
                            item_value = indicator['value'].lower().replace(path, p_item)
                            item_id = str(indicator['id']) + "_" + str(subindicator)
                            string_buffer.write('          ${} = "{}" {}\n'.format(item_id, format_yara_string(item_value), string_modifiers[indicator['type'].lower()]))
                            subindicator += 1

                if subindicator == 0:
                    string_buffer.write('          $sip_{} = "{}" {}\n'.format(item_id, format_yara_string(item_value), string_modifiers[indicator['type'].lower()]))

            elif indicator['type'] == get_indicator_type_mapping(saq.intel.I_REGISTRY_KEY):
                for reg in ['hkcu\\', 'hklm\\', 'hkc\\', 'hku\\', 'hkcr\\']:
                    item_value = item_value.lower().replace(reg, "") # remove the front end of the indicator if it matches our special case
                string_buffer.write('          $sip_{} = "{}" {}\n'.format(item_id, format_yara_string(item_value), string_modifiers[indicator['type'].lower()]))
            elif indicator['type'] == get_indicator_type_mapping(saq.intel.I_HEX_STRING):
                # hex strings are passed as is and are expected to be in yara hex string format
                string_buffer.write('          $sip_{} = {{ {} }}\n'.format(item_id, item_value))
            else:
                string_buffer.write('          $sip_{} = "{}" {}\n'.format(item_id, format_yara_string(item_value), string_modifiers[indicator['type'].lower()]))

            # make sure this individual string compiles into a yara rule
            test_rule = template.replace('TEMPLATE_STRINGS', string_buffer.getvalue())
            try:
                yara.compile(source=test_rule)
            except Exception as e:
                logging.error("indicator type {} id {} value {} invalid for yara".format(indicator_type, item_id, item_value))
                print(test_rule)
                skip_count += 1
                continue

            count += 1
            string_data.write(string_buffer.getvalue())

        if count == 0:
            logging.info("no indicators of type {} available for export".format(indicator_type))
            continue

        # make sure the whole rule compiles
        rule = template.replace('TEMPLATE_STRINGS', string_data.getvalue()) 
        try:
            yara.compile(source=rule)
        except Exception as e:
            logging.error("unable to compile rules for {}: {}".format(indicator_type, e))
            print(rule)
            continue

        # make sure something changed
        output_file = get_yara_filename(indicator_type)
        if os.path.exists(output_file):
            with open(output_file, 'r') as fp:
                existing_rule = fp.read()

            if existing_rule == rule:
                logging.info("no changes detected for {}".format(indicator_type))
                continue

        with open(output_file, 'w') as fp:
            fp.write(rule)

        logging.info("exported {} skipped {} indicators of type {} to {}".format(count, skip_count, indicator_type, output_file))
        if count == 0:
            logging.warning("no strings were exported for {}, removing {}".format(indicator_type, output_file))
            try:
                os.remove(output_file)
            except Exception as e:
                logging.error("unable to remove {}: {}".format(output_file, e))

    sys.exit(0)

export_sip_yara_rules_parser = subparsers.add_parser('export-sip-yara-rules',
    help="Exports your SIP database as yara rules to be used by ACE.")
export_sip_yara_rules_parser.add_argument('dir',
    help="The path to the directory to export the yara rules into.")
export_sip_yara_rules_parser.set_defaults(func=export_sip_yara_rules)

def export_crits_yara_rules(args):

    import io

    from collections import defaultdict

    import saq
    import saq.intel

    from saq.crits import get_indicator_type_mapping 

    from pymongo import MongoClient

    import yara
    yara.set_config(max_strings_per_rule=30720)

    # make sure target directory exists
    if not os.path.isdir(args.dir):
        try:
            os.makedirs(args.dir)
        except Exception as e:
            logging.error("unable to create directory {}: {}".format(args.dir, e))
            sys.exit(1)

    # load the mapping from indicator type to the string modifiers to use
    string_modifiers = defaultdict(lambda: saq.CONFIG['crits_yara_export_string_modifiers']['default'])
    for indicator_type in saq.CONFIG['crits_yara_export_string_modifiers']:
        if indicator_type == 'default':
            continue

        string_modifiers[indicator_type] = saq.CONFIG['crits_yara_export_string_modifiers'][indicator_type]
        logging.debug("using string modifiers {} for {}".format(string_modifiers[indicator_type], indicator_type))

    # what is the minimum length an indicator value can be to be put into a yara rule?
    export_minimum_length = saq.CONFIG['crits_yara_export'].getint('export_minimum_length')

    def format_yara_string(s):
        return s.replace('\\', '\\\\').replace('"','\\"').replace("\n","")

    def get_yara_filename(indicator_type):
        return os.path.join(args.dir, "CRITS_{}.yar".format(indicator_type))

    with MongoClient(saq.CONFIG['crits']['mongodb_uri']) as connection:
        db = connection['crits']
        
        indicator_types = [_.strip() for _ in saq.CONFIG['crits_yara_export']['export_list'].split(',')]
        included_sources = [_.strip() for _ in saq.CONFIG['crits_yara_export']['export_sources_include'].split(',')]
        excluded_sources = [_.strip() for _ in saq.CONFIG['crits_yara_export']['export_sources_exclude'].split(',')]

        sources = []
        for source in db.indicators.distinct('source.name'):
            if included_sources and source not in included_sources:
                continue

            if source in excluded_sources:
                continue

            sources.append(source)

        for indicator_type in indicator_types:
            # what is the crits equivalent of this indicator type?
            logging.debug("exporting indicator type {}".format(indicator_type))
          
            # does a template file exists for this indicator type?
            template_path = os.path.join(saq.CONFIG['crits_yara_export']['export_template_dir'], '{}.template'.format(indicator_type))
            if not os.path.exists(template_path):
                template_path = os.path.join(saq.CONFIG['crits_yara_export']['export_template_dir'], 'default.template')

            logging.info("using template {} for {}".format(template_path, indicator_type))

            # load the template we're going to be using
            with open(template_path, 'r') as fp:
                template = fp.read()

            template = template.replace('TEMPLATE_RULE_NAME', 'CRITS_{}'.format(re.sub(r'[^a-zA-Z0-9_]', '', indicator_type)))
            
            special_paths = { "%temp%":[ "\\windows\\temp", "\\temp", "\\appdata\\local\\temp", "\\local settings\\temp", "\\locals~1\\temp" ],
                              "%appdata%": [ "\\application data", "\\appdata\\roaming" ],
                              "%programdata%": [ "\\programdata", "\\documents and settings\\all users" ],
                              "%programfiles%": [ "\\program files", "\\program files (x86)" ],
                              "%systemdrive%": [ "" ],
                              "%system%": [ "\\windows\\system32", "\\windows\\system" ] }

            string_data = io.StringIO()
            count = 0
            skip_count = 0

            clause = {"status": "Analyzed", "type": get_indicator_type_mapping(indicator_type)}
            if sources:
                clause['source.name'] = { '$in': sources }

            for item in db.indicators.find(clause):
                string_buffer = io.StringIO()
                item_id = item['_id']
                item_value = item['value']

                # is this string too small?
                if len(item_value) < export_minimum_length:
                    skip_count += 1
                    continue
                    
                # do we need to reformat this string?
                if item['type'] == get_indicator_type_mapping(saq.intel.I_FILE_PATH):
                    subindicator = 0
                    for path in special_paths:
                        if path.lower() in item['value'].lower():
                            for p_item in special_paths[path]:
                                item_value = item['value'].lower().replace(path, p_item)
                                item_id = str(item['_id']) + "_" + str(subindicator)
                                string_buffer.write('          ${} = "{}" {}\n'.format(item_id, format_yara_string(item_value), string_modifiers[item['type'].lower()]))
                                subindicator += 1

                    if subindicator == 0:
                        string_buffer.write('          ${} = "{}" {}\n'.format(item_id, format_yara_string(item_value), string_modifiers[item['type'].lower()]))

                elif item['type'] == get_indicator_type_mapping(saq.intel.I_REGISTRY_KEY):
                    for reg in ['hkcu\\', 'hklm\\', 'hkc\\', 'hku\\', 'hkcr\\']:
                        item_value = item_value.lower().replace(reg, "") # remove the front end of the indicator if it matches our special case
                    string_buffer.write('          ${} = "{}" {}\n'.format(item_id, format_yara_string(item_value), string_modifiers[item['type'].lower()]))
                elif item['type'] == get_indicator_type_mapping(saq.intel.I_HEX_STRING):
                    # hex strings are passed as is and are expected to be in yara hex string format
                    string_buffer.write('          ${} = {{ {} }}\n'.format(item_id, item_value))
                else:
                    string_buffer.write('          ${} = "{}" {}\n'.format(item_id, format_yara_string(item_value), string_modifiers[item['type'].lower()]))

                # make sure this individual string compiles into a yara rule
                test_rule = template.replace('TEMPLATE_STRINGS', string_buffer.getvalue())
                try:
                    yara.compile(source=test_rule)
                except Exception as e:
                    logging.error("indicator type {} id {} value {} invalid for yara".format(indicator_type, item_id, item_value))
                    print(test_rule)
                    skip_count += 1
                    continue

                count += 1
                string_data.write(string_buffer.getvalue())

            if count == 0:
                logging.info("no indicators of type {} available for export".format(indicator_type))
                continue

            # make sure the whole rule compiles
            rule = template.replace('TEMPLATE_STRINGS', string_data.getvalue()) 
            try:
                yara.compile(source=rule)
            except Exception as e:
                logging.error("unable to compile rules for {}: {}".format(indicator_type, e))
                print(rule)
                continue

            # make sure something changed
            output_file = get_yara_filename(indicator_type)
            if os.path.exists(output_file):
                with open(output_file, 'r') as fp:
                    existing_rule = fp.read()

                if existing_rule == rule:
                    logging.info("no changes detected for {}".format(indicator_type))
                    continue

            with open(output_file, 'w') as fp:
                fp.write(rule)

            logging.info("exported {} skipped {} indicators of type {} to {}".format(count, skip_count, indicator_type, output_file))
            if count == 0:
                logging.warning("no strings were exported for {}, removing {}".format(indicator_type, output_file))
                try:
                    os.remove(output_file)
                except Exception as e:
                    logging.error("unable to remove {}: {}".format(output_file, e))
    
    sys.exit(0)

export_crits_yara_rules_parser = subparsers.add_parser('export-crits-yara-rules',
    help="Exports your CRITS database as yara rules to be used by ACE.")
export_crits_yara_rules_parser.add_argument('dir',
    help="The path to the directory to export the yara rules into.")
export_crits_yara_rules_parser.set_defaults(func=export_crits_yara_rules)

def list_observables(args):
    from saq.constants import VALID_OBSERVABLE_TYPES, OBSERVABLE_DESCRIPTIONS
    print() 
    print("SAQ Valid Observable Types")
    print()
    for o_type in VALID_OBSERVABLE_TYPES:
        print('{0}{1}'.format(o_type.ljust(20), OBSERVABLE_DESCRIPTIONS[o_type]))
    print()

list_observables_parsers = subparsers.add_parser('list-observables',
    help="List available observable types and their descriptions.")
list_observables_parsers.set_defaults(func=list_observables)

# XXX get rid of this
def crits_activity_sync(args):
    import saq
    from saq.constants import DISPOSITION_IGNORE
    from saq.crits import sync_crits_activity
    from saq.database import Alert, DatabaseSession
    session = DatabaseSession()
    alerts = session.query(Alert).filter(
        Alert.location == saq.CONFIG['global']['node'],
        Alert.disposition != DISPOSITION_IGNORE,
        Alert.disposition_time >= datetime.datetime.now() + datetime.timedelta(days=-args.numdays))

    session.commit()

    for alert in alerts:
        try:
            alert.load()
        except Exception as e:
            logging.error("unable to load alert {}: {}".format(alert, e))
            continue

        if "indicator" in alert.observable_types:
            sync_crits_activity(alert)

crits_activity_sync_parsers = subparsers.add_parser('crits-activity-sync',
    help="Sync activity from this alert to CRITS alert activity")
crits_activity_sync_parsers.add_argument('-d','--number-days-to-sync',dest='numdays',default=1,type=int,
    help="Number of days in the past to sync activity to CRITS, default is previous day")
crits_activity_sync_parsers.set_defaults(func=crits_activity_sync)

def interactive_console(args):
    import pdb
    pdb.set_trace()
    sys.exit(0)

interactive_console_parser = subparsers.add_parser('interactive-console',
    help="Drop into the REPL after initialization of ACE.")
interactive_console_parser.set_defaults(func=interactive_console)

def cleanup_alerts(args):
    """Performs system maintenance.  This is meant to be called from a cron job."""
    from saq.util.maintenance import cleanup_alerts
    cleanup_alerts(fp_days_old=args.fp_days_old, 
                   ignore_days_old=args.ignore_days_old,
                   dry_run=args.dry_run)
    sys.exit(0)

cleanup_alerts_parsers = subparsers.add_parser('cleanup-alerts',
    help="Removes alerts dispositioned as ignore or false positive and older than some amount of time.")
cleanup_alerts_parsers.add_argument('--dry-run', required=False, dest='dry_run', default=False, action='store_true',
    help="Just report how many would be deleted and archived.")
cleanup_alerts_parsers.add_argument('--fp-days-old', type=int, required=False, dest='fp_days_old', default=None, action='store',
    help='Specify how many days old an alert dispositioned as FALSE_POSITIVE should be for it to be archived.')
cleanup_alerts_parsers.add_argument('--ignore-days-old', type=int, required=False, dest='ignore_days_old', default=None, action='store',
    help='Specify how many days old an alert dispositioned as IGNORE should be for it to be deleted.')
#cleanup_alerts_parsers.add_argument('--force-delete', required=False, dest='force_delete', default=None, action='store_true',
    #help='force delete fp alerts instead of archiving them')
cleanup_alerts_parsers.set_defaults(func=cleanup_alerts)

def cleanup_email_archive(args):
    from saq.email import maintain_archive
    maintain_archive(verbose=True)
    sys.exit(0)

cleanup_email_archive_parser = subparsers.add_parser('cleanup-email-archive',
    help="Removes emails from the archive that have expired.")
cleanup_email_archive_parser.set_defaults(func=cleanup_email_archive)

def cleanup_cloudphish(args):
    from saq.database import get_db_connection

    delete_counter = 0
    url_counter = 0

    with get_db_connection() as db:
        c = db.cursor()

        if args.sha256_url:
            c.execute("""
SELECT 
    LOWER(HEX(cloudphish_analysis_results.sha256_url)),
    LOWER(HEX(cloudphish_analysis_results.sha256_content))
FROM 
    cloudphish_analysis_results  
WHERE
    cloudphish_analysis_results.sha256_url = UNHEX(%s)
""", ( args.sha256_url, ))

        else:
            c.execute("""
SELECT 
    LOWER(HEX(cloudphish_analysis_results.sha256_url)),
    LOWER(HEX(cloudphish_analysis_results.sha256_content))
FROM 
    cloudphish_url_lookup JOIN cloudphish_analysis_results ON cloudphish_url_lookup.sha256_url = cloudphish_analysis_results.sha256_url 
WHERE
    cloudphish_analysis_results.result IN ( 'UNKNOWN', 'ERROR', 'CLEAR', 'PASS' )
    AND status = 'ANALYZED'
    AND cloudphish_url_lookup.last_lookup < NOW() - INTERVAL %s DAY
""", ( args.cleanup_days, ))

        results = c.fetchall()
        db.commit()

        if args.dry_run:
            for sha256_url, sha256_content in results:
                print("Would delete - sha256_url:{} sha256_content:{}".format(sha256_url, sha256_content))
            sys.exit(0)

        logging.info(f"clearing {len(results)} cloudphish urls")
        for sha256_url, sha256_content in results:
            if sha256_content:
                target_file = os.path.join(saq.DATA_DIR, saq.CONFIG['cloudphish']['cache_dir'], 
                                           sha256_content[:2], sha256_content)

                if os.path.exists(target_file):
                    try:
                        os.remove(target_file)
                        delete_counter += 1
                    except Exception as e:
                        logging.error(f"unable to delete {target_file}: {e}")

        for sha256_url, sha256_content in results:
            c.execute("DELETE FROM cloudphish_analysis_results WHERE sha256_url = UNHEX(%s)", (sha256_url,))
            url_counter += 1
            if url_counter % 100 == 0:
                db.commit()

    logging.info(f"removed {url_counter} urls and deleted {delete_counter} files")
    sys.exit(0)

cleanup_cloudphish_parser = subparsers.add_parser('cleanup-cloudphish',
    help="Cleans out old urls and content cached by cloudphish requests.")
cleanup_cloudphish_parser.add_argument('--sha256-url', required=False, dest='sha256_url',
    default=False, action='store', help="the sha256 of the url that you want to cleanup for")
cleanup_cloudphish_parser.add_argument('--dry-run', required=False, dest='dry_run',
    default=False, action='store_true', help="Just return the results that would be purged")
cleanup_cloudphish_parser.add_argument('--cleanup-days', type=int, default=7,
    help="Select URLs that have not been requested in X days (defaults to 7 days.)")
cleanup_cloudphish_parser.set_defaults(func=cleanup_cloudphish)

def display_alert(args):
    from saq.analysis import RootAnalysis
    
    alert = RootAnalysis()
    alert.storage_dir = args.dir
    try:
        alert.load()
    except Exception as e:
        logging.error("unable to load alert from {}: {}".format(args.dir, str(e)))
        traceback.print_exc()
        sys.exit(1)

    display_analysis(alert)
    sys.exit(0)

display_alert_parser = subparsers.add_parser('display-alert',
    help="Displays the results of the analysis for a given alert.")
display_alert_parser.add_argument('dir', 
    help="The directory of the alert to display")
display_alert_parser.set_defaults(func=display_alert)

# ============================================================================
# company management
#

def list_companies(args):
    from saq.database import get_db_connection

    with get_db_connection() as db:
        c = db.cursor()
        c.execute("SELECT id, name FROM company ORDER BY name")
        print()
        print("ID\tNAME")
        for company_id, company_name in c:
            print("{}\t{}".format(company_id, company_name))

    print()
    print("use ./saq add-company and ./saq delete-company to manage companies")
    print()
    sys.exit(0)

list_companies_parser = subparsers.add_parser('list-companies',
    help="Lists the available companies and their IDs.")
list_companies_parser.set_defaults(func=list_companies)

def add_company(args):
    from saq.database import get_db_connection

    with get_db_connection() as db:
        c = db.cursor()
        c.execute("INSERT INTO company ( `id`, `name` ) VALUES ( %s, %s )", (args.company_id, args.company_name))
        db.commit()
        logging.info("company added")

    sys.exit(0)

add_companies_parser = subparsers.add_parser('add-company',
    help="Adds a new company entry.")
add_companies_parser.add_argument('company_id', type=int, help="The ID of the new company (a number that is not already being used as an ID.)")
add_companies_parser.add_argument('company_name', help="The name of the new company.")
add_companies_parser.set_defaults(func=add_company)

def delete_company(args):
    import saq
    from saq.constants import INSTANCE_TYPE_PRODUCTION
    from saq.database import get_db_connection

    print("***************************************************************")
    print("Deleting a company will delete all associated EVENTS and ALERTS.")
    confirm = input("Are you SURE? (Y/n)")
    if confirm != 'Y':
        print("Action not taken.")
        sys.exit(0)

    confirm = input("Are you DAMN SURE? Seriously. If you're wrong it will be a disaster. (Y/n)")
    if confirm != 'Y':
        print("Action not taken.")
        sys.exit(0)

    # see if we are on the production system
    if saq.INSTANCE_TYPE == INSTANCE_TYPE_PRODUCTION:
        confirm = input("You are in a PRODUCTION SERVER. Are you SURE you know what you are doing? Type YES if you are sure.")
        if confirm != 'YES':
            print("Action not taken. Pay attention to what you're doing please.")
            sys.exit(0)

    with get_db_connection() as db:
        c = db.cursor()
        c.execute("DELETE FROM company WHERE `name` = %s", (args.company_name,))
        db.commit()
        logging.info("company deleted")

    sys.exit(0)

delete_companies_parser = subparsers.add_parser('delete-company',
    help="Deletes a given company and all associated events and alerts.")
delete_companies_parser.add_argument('company_name', help="The name of the company to delete.")
delete_companies_parser.set_defaults(func=delete_company)

# ============================================================================
# testing utilities
#

def check_url(args):
    from saq.crawlphish import CrawlphishURLFilter

    cw = CrawlphishURLFilter()
    cw.load()

    result = cw.filter(args.url)
    print(f"filtered: {result.filtered} reason: {result.reason}")

    #if cw.is_filtered(args.url):
        #print("url filtered (reason {})".format(cw.reason))
    #else:
        #print("url OK (reason {})".format(cw.reason))

    sys.exit(0)

check_url_parser = subparsers.add_parser('check-url',
    help="Check the given URL against crawlphish URL selection logic.")
check_url_parser.add_argument('url', help="The URL to check.")
check_url_parser.set_defaults(func=check_url)

def test_daemonize(args):
    if args.kill_daemon:
        kill_daemon('test-daemon')

    if args.daemon:
        daemonize('test-daemon')

    while True:
        logging.info('running')
        time.sleep(1)

test_daemonize_parsers = subparsers.add_parser('test-daemonize',
    help="Test daemonization.")
test_daemonize_parsers.set_defaults(func=test_daemonize)

def test_proxies(args):
    import requests
    import saq
    proxy_configs = saq.OTHER_PROXIES.copy()
    proxy_configs['GLOBAL'] = saq.PROXIES

    for proxy_name, proxy_config in proxy_configs.items():
        print("testing proxy {} ({})".format(proxy_name, proxy_config))
        session = requests.session()
        session.proxies = proxy_config
        response = session.request('GET', args.url,
                                   timeout=20,
                                   allow_redirects=True,
                                   verify=False)

        print("result: ({}) - {}".format(response.status_code, response.reason))

    sys.exit(0)

test_proxies_parser = subparsers.add_parser('test-proxies',
    help="Test the configured proxies to make sure ACE can use them.")
test_proxies_parser.add_argument('url', help="A sample URL to attempt to download through each proxy.")
test_proxies_parser.set_defaults(func=test_proxies)

def test_process_server(args):
    from saq.process_server import initialize_process_server, Popen, PIPE, TimeoutExpired
    initialize_process_server()

    try:
        p = Popen(['./test_sleep',], stdout=PIPE, stderr=PIPE)
        _stdout, _stderr = p.communicate(timeout=2)
        print("_stdout = {}".format(_stdout))
        print("_stderr = {}".format(_stderr))
        print("returncode = {}".format(p.returncode))
    except TimeoutExpired as e:
        print("timeout OK")

    p = Popen(['cat'], stdin=PIPE, stdout=PIPE, universal_newlines=True)
    p.stdin.write('Hello world!')
    _stdout, _stderr = p.communicate()
    print("_stdout = {}".format(_stdout))
    print("_stderr = {}".format(_stderr))

    p = Popen(['cat'], stdin=PIPE, stdout=PIPE)
    p.stdin.write(b'Hello world!')
    _stdout, _stderr = p.communicate()
    print("_stdout = {}".format(_stdout))
    print("_stderr = {}".format(_stderr))

    p = Popen(['/bin/false'])
    p.wait(timeout=5)
    print("rcode = {}".format(p.returncode))

    with open('large_file_copy.stderr', 'wb') as fp_stderr:
        with open('large_file_copy', 'wb') as fp:
            p = Popen(['bash', '-c', 'cat large_file 1>&2'], stdout=fp, stderr=fp_stderr)
            p.wait()

test_process_server_parser = subparsers.add_parser('test-process-server',
    help="Test process server.")
test_process_server_parser.set_defaults(func=test_process_server)

# XXX replace this with calls to the engine code
def verify_modules(args):
    """Executes verify_environment() on all modules that are enabled."""
    # we run the same code the engines run to load the modules
    # COPY-PASTA!!
    import importlib

    for section in saq.CONFIG.sections():
        if not section.startswith('analysis_module_'):
            continue

        # is this module disabled globally?
        # modules that are disable globally are not used anywhere
        if not saq.CONFIG.getboolean(section, 'enabled'):
            logging.warning("analysis module {} disabled (globally)".format(section))
            continue

        logging.debug("verifying analysis module from {}".format(section))
        module_name = saq.CONFIG.get(section, 'module')
        try:
            _module = importlib.import_module(module_name)
        except Exception as e:
            logging.error("unable to import module {}".format(module_name, e))
            traceback.print_exc()
            continue

        class_name = saq.CONFIG.get(section, 'class')
        try:
            module_class = getattr(_module, class_name)
        except AttributeError as e:
            logging.error("class {} does not exist in module {} in analysis module {}".format(
                          class_name, module_name, section))
            traceback.print_exc()
            continue

        try:
            analysis_module = module_class(section)
        except Exception as e:
            logging.error("unable to load analysis module {}: {}".format(section, e))
            traceback.print_exc()
            continue

        # make sure the module has everything it needs
        try:
            analysis_module.verify_environment()
        except Exception as e:
            logging.error("analysis module {} failed environment verification: {}".format(analysis_module, e))
            traceback.print_exc()
            continue

        logging.info("analysis module {} verification OK".format(section))
    
verify_modules_parsers = subparsers.add_parser('verify-modules',
    help="Executes verify_environment() on all modules that are enabled.")
verify_modules_parsers.set_defaults(func=verify_modules)

def set_encryption_password(args):
    from saq.crypto import set_encryption_password
    while True:
        password = getpass.getpass("Enter the decryption password:")
        password_2 = getpass.getpass("Re-enter the decryption password for verification:")

        if password != password_2:
            logging.error("passwords do not match")
            continue

        break

    set_encryption_password(password)
    sys.exit(0)

set_encryption_password_parser = subparsers.add_parser('set-encryption-password',
    help="Sets the password used to encrypt and decrypt archived emails.")
set_encryption_password_parser.set_defaults(func=set_encryption_password)

def test_encryption_password(args):
    from saq.crypto import test_encryption_password
    password = getpass.getpass("Enter the decryption password:")
    if test_encryption_password(password):
        logging.info("password OK")
    else:
        logging.error("invalid password")

    sys.exit(0)

test_encryption_password_parser = subparsers.add_parser('test-encryption-password',
    help="Tests the given password to see if it matches what is currently set as the encryption password.")
test_encryption_password_parser.set_defaults(func=test_encryption_password)

def encrypt_file(args):
    from saq.crypto import encrypt
    encrypt(args.source_path, args.target_path)
    sys.exit(0)

encrypt_file_parser = subparsers.add_parser('encrypt-file',
    help="Encrypts the given file with the password set with set-encryption-password.")
encrypt_file_parser.add_argument('source_path', help="The file to encrypt from.")
encrypt_file_parser.add_argument('target_path', help="The file to saved the decrypted data to.")
encrypt_file_parser.set_defaults(func=encrypt_file)

def decrypt_file(args):
    from saq.crypto import decrypt
    decrypt(args.source_path, args.target_path)
    sys.exit(0)

decrypt_file_parser = subparsers.add_parser('decrypt-file',
    help="Decrypts the given file with the password set with set-decryption-password.")
decrypt_file_parser.add_argument('source_path', help="The file to decrypt from.")
decrypt_file_parser.add_argument('target_path', help="The file to saved the decrypted data to.")
decrypt_file_parser.set_defaults(func=decrypt_file)

def add_bro_http_whitelist(args):
    import saq

    if not re.match(r'^(\d{1,3}\.){3}\d{1,3}/\d+$', args.cidr):
        logging.error("invalid CIDR (use a.b.c.d/e format)")
        sys.exit(1)

    if not args.description:
        logging.error("a brief description is required")
        sys.exit(1)

    with open(os.path.join(saq.SAQ_HOME, 'bro', 'http.whitelist'), 'a') as fp:
        fp.write('{}\t{}\n'.format(args.cidr, args.description))

    sys.exit(0)

add_bro_http_whitelist_parser = subparsers.add_parser('add-bro-http-whitelist',
    help="Adds the given CIDR and description as a whitelist item to the bro HTTP whitelist.")
add_bro_http_whitelist_parser.add_argument('cidr', help="The network CIDR to whitelist.")
add_bro_http_whitelist_parser.add_argument('description', help="A description of the whitelisted item.")
add_bro_http_whitelist_parser.set_defaults(func=add_bro_http_whitelist)

def remove_bro_http_whitelist(args):
    removed_line = False
    src_path = os.path.join(saq.SAQ_HOME, 'bro', 'http.whitelist')
    tmp_path = os.path.join(saq.SAQ_HOME, 'bro', 'http.whitelist.tmp')

    with open(src_path, 'r') as fp_in:
        with open(tmp_path, 'a') as fp_out:
            for line in fp_in:
                if line.startswith(args.cidr):
                    logging.info("removed {}".format(line.strip()))
                    removed_line = True
                else:
                    fp_out.write(line)

    if removed_line:
        shutil.copy(tmp_path, src_path)
    
    os.remove(tmp_path)
    sys.exit(0)

remove_bro_http_whitelist_parser = subparsers.add_parser('remove-bro-http-whitelist',
    help="Adds the given CIDR and description as a whitelist item to the bro HTTP whitelist.")
remove_bro_http_whitelist_parser.add_argument('cidr', help="The network CIDR to remove from the whitelist.")
remove_bro_http_whitelist_parser.set_defaults(func=remove_bro_http_whitelist)

def list_available_modules(args):
    from saq.engine import load_module

    result = {}

    for section in saq.CONFIG.keys():
        if not section.startswith('analysis_module_'):
            continue

        if section in saq.CONFIG['disabled_modules'] and saq.CONFIG['disabled_modules'].getboolean(section):
            continue

        result[section] = load_module(section)

    for key in sorted(result.keys()):
        print('{: <35}{}'.format(key[len('analysis_module_'):], result[key].__doc__ if result[key].__doc__ is not None else ''))

    sys.exit(0)

list_available_modules_parser = subparsers.add_parser('list-available-modules',
    help="Lists the modules available in ACE.")
list_available_modules_parser.set_defaults(func=list_available_modules)

def config(args):
    import saq

    def matches_param(s, k=None):
        if not args.settings:
            return True

        for spec in args.settings:
            section_spec, key_spec = spec.split('.')
            #print("testing {} == {} {} == {}".format(section_spec, s, key_spec, k))
            if ( section_spec == '*' or section_spec == s ) and ( k is None or key_spec == '*' or key_spec == k ):
                return True

        return False

    for section in list(saq.CONFIG.keys()):
        if section == 'DEFAULT':
            continue

        if not matches_param(section):
            continue

        if not args.value_only:
            print('[{}]'.format(section))
        
        for key in saq.CONFIG[section].keys():
            if not matches_param(section, key):
                continue

            if args.value_only:
                print(saq.CONFIG[section][key])
            else:
                print('{} = {}'.format(key, saq.CONFIG[section][key]))

        if not args.value_only:
            print()

    sys.exit(0)

config_parser = subparsers.add_parser('config',
    help="Queries the ACE configuration.")
config_parser.add_argument('-v', '--value', action='store_true', default=False, dest='value_only',
    help="Just print the values of the selected configuration items.")
config_parser.add_argument('settings', nargs="*",
    help="Zero or more configuration items to display in the format section.key.")
config_parser.set_defaults(func=config)

def display_workload(args):
    from saq.database import get_db_connection

    with get_db_connection() as db:
        c = db.cursor()
        c.execute("SELECT analysis_mode, COUNT(*) FROM workload WHERE company_id = %s GROUP BY analysis_mode ORDER BY analysis_mode", (saq.COMPANY_ID,))
        print(f" -- WORKLOAD ({saq.COMPANY_NAME}) --")
        print("{: <15}{}".format('MODE', 'COUNT'))
        for analysis_mode, count in c:
            print("{: <15}{}".format(analysis_mode, count))

        db.commit()
        print()
        print(f" -- DELAYED WORKLOAD ({saq.SAQ_NODE})--")
        c.execute("SELECT storage_dir, analysis_module, COUNT(*) FROM delayed_analysis JOIN nodes ON delayed_analysis.node_id = nodes.id WHERE nodes.name = %s GROUP BY storage_dir, analysis_module", (saq.SAQ_NODE,))
        print("{: <36} {: <20} {}".format('UUID', 'MODULE', 'COUNT'))
        for storage_dir, analysis_module, count in c:
            storage_dir = os.path.basename(storage_dir)
            analysis_module = analysis_module[len('analysis_module_'):]
            print("{: <36} {: <20} {}".format(storage_dir, analysis_module, count))

        db.commit()
        print()
        print(f" -- LOCKS ({saq.SAQ_NODE}) --")
        c.execute("SELECT uuid, lock_uuid, lock_time, lock_owner FROM locks WHERE lock_owner LIKE CONCAT(%s, '-%%') ORDER BY lock_time", (saq.SAQ_NODE,))
        print("{: <36} {: <36} {: <19} {}".format('UUID', 'LOCK', 'TIME', 'OWNER'))
        for _uuid, lock_uuid, lock_time, lock_owner in c:
            print("{: <36} {: <36} {: <19} {}".format(_uuid, lock_uuid, str(lock_time), lock_owner))
        db.commit()

    sys.exit(0)

display_workload_parser = subparsers.add_parser('display-workload',
    help="Displays the current ACE workload.")
display_workload_parser.set_defaults(func=display_workload)

if __name__ == '__main__':

    # there is no reason to run anything as root
    if os.geteuid() == 0:
        print("do not run ace as root please")
        sys.exit(1)

    # check the current version of python
    if sys.version_info < (3, 6):
        print("you need at least python version 3.6.* to run ACE")
        sys.exit(1)

    # was this command called as a symlink?
    if os.path.basename(sys.argv[0]) == 'correlate':
        sys.argv.insert(1, 'correlate')

    # parse the command line arguments
    args = parser.parse_args()

    # where is ACE?
    saq_home = '/opt/ace' # default installation directory
    if 'SAQ_HOME' in os.environ:
        saq_home = os.environ['SAQ_HOME'] # should be in this env variable
    if args.saq_home is not None:
        saq_home = args.saq_home # otherwise pull it from the command line

    # XXX I would like to get rid of this
    # adjust search path
    if args.saq_home is not None:
        sys.path.append(os.path.join(args.saq_home, 'lib'))
    elif 'SAQ_HOME' in os.environ:
        sys.path.append(os.path.join(os.environ['SAQ_HOME'], 'lib'))

    # initialize saq
    import saq
    saq.initialize(saq_home=saq_home, config_paths=[], logging_config_path=None, args=args, relative_dir=args.relative_dir)

    # are we prompting for the decryption password?
    if args.provide_decryption_password:
        from saq.crypto import test_encryption_password, get_aes_key

        while True:
            password = getpass.getpass("Enter the decryption password:")
            # first check to see if the given password matches the sha256 hash
            if not test_encryption_password(password):
                logging.error("invalid encryption password")
                continue

            break

        saq.ENCRYPTION_PASSWORD = get_aes_key(password)

    if args.debug_on_error:
        def info(type, value, tb):
           if hasattr(sys, 'ps1') or not sys.stderr.isatty() or type != AssertionError:
              # we are in interactive mode or we don't have a tty-like
              # device, so we call the default hook
              sys.__excepthook__(type, value, tb)
           else:
              import traceback, pdb
              # we are NOT in interactive mode, print the exception...
              traceback.print_exception(type, value, tb)
              print
              # ...then start the debugger in post-mortem mode.
              pdb.pm()

        sys.excepthook = info

    # call the handler for the given command
    args.func(args)