ace

#!/usr/bin/env python3
# PYTHON_ARGCOMPLETE_OK

import argcomplete
import argparse
import atexit
import collections
import datetime
import fnmatch
import getpass
import importlib
import io
import json
import logging
import operator
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="Analysis Correlation Engine")
parser.add_argument('--saq-home', required=False, dest='saq_home', default=None,
    help="Sets the base directory of ACE.")
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.
         This option can be specified multiple times and each file is loaded in order.""")
#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. Read README.CRYPTO for details.")
parser.add_argument('-P', '--set-decryption-password', dest='set_decryption_password', default=None,
    help="Provide the decryption password on the command line. Not secure at all. Don't do it.")
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.debug("removing proxy setting {}".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)

    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))

# ============================================================================
# hunting utilities
#

def execute_hunt(args):
    import ace_api
    import pytz
    from saq.constants import event_time_format_tz
    from saq.collectors.hunter import HunterCollector
    from saq.collectors.query_hunter import QueryHunt
    collector = HunterCollector()
    collector.load_hunt_managers()
    
    hunt_type, hunt_name = args.hunt.split(':', 1)
    if hunt_type not in collector.hunt_managers:
        logging.error(f"invalid hunt type {hunt_type}")
        sys.exit(1)

    collector.hunt_managers[hunt_type].load_hunts_from_config()
    hunts = collector.hunt_managers[hunt_type].get_hunts(lambda hunt: hunt_name in hunt.ini_path)
    if hunts is None:
        logging.error(f"unknown hunt {hunt_name} for type {hunt_type}")
        sys.exit(1)
    hunt = hunts[0]
    if len(hunts) > 1:
        logging.warning(f"got {len(hunts)} hunts matching '{hunt_name}' for type {hunt_type}")
        for _hunt in hunts:
            _ini_name = os.path.basename(_hunt.ini_path).split('.', 1)[0]
            logging.warning(f" -> {_hunt.name} @ {_hunt.ini_path} - {_ini_name}")
            if hunt_name == _ini_name:
                hunt = _hunt
        logging.warning(f"going with {hunt.name} @ {hunt.ini_path}...")

    # set the Hunt to manual so we don't record the execution timestamps
    hunt.manual_hunt = True
    exec_kwargs = {}

    if isinstance(hunt, QueryHunt):
        start_time = datetime.datetime.strptime(args.start_time, '%m/%d/%Y:%H:%M:%S')
        end_time = datetime.datetime.strptime(args.end_time, '%m/%d/%Y:%H:%M:%S')
        if args.timezone is not None:
            start_time = pytz.timezone(args.timezone).localize(start_time)
            end_time = pytz.timezone(args.timezone).localize(end_time)
        else:
            start_time = pytz.utc.localize(start_time)
            end_time = pytz.utc.localize(end_time)

        exec_kwargs['start_time'] = start_time
        exec_kwargs['end_time'] = end_time

        hunt.query_result_file = args.query_result_file

    if args.json_dir is not None:
        os.makedirs(args.json_dir, exist_ok=True)

    json_dir_index = 0
    for submission in hunt.execute(**exec_kwargs):
        print(submission.description)
        if args.details:
            for o in submission.observables:
                output = f"\t(*) {o['type']} - {o['value']}"
                if 'time' in o:
                    output += " - {}".format(o['time'].strftime(event_time_format_tz))
                if 'tags' in o:
                    output += " tags [{}]".format(','.join(o['tags']))
                if 'directives' in o:
                    output += " direc [{}]".format(','.join(o['directives']))
                # TODO the other stuff
                print(output)

            for t in submission.tags:
                print(f"\t(+) {t}")
            
        if args.events:
            print("BEGIN EVENTS")
            for event in submission.details:
                print(event)
            print("END EVENTS")

        if args.json_dir is not None:
            buffer = []
            for event in submission.details:
                buffer.append(event)

            target_json_file = os.path.join(args.json_dir, '{}.json'.format(str(json_dir_index)))
            with open(target_json_file, 'w') as fp:
                json.dump(buffer, fp)

            with open(os.path.join(args.json_dir, 'manifest'), 'a') as fp:
                fp.write(f'{json_dir_index} = {submission.description}\n')

            json_dir_index += 1

        if args.submit_alerts is not None:
            result = ace_api.submit(
                submission.description,
                remote_host=args.submit_alerts,
                ssl_verification=saq.CONFIG['SSL']['ca_chain_path'],
                analysis_mode=submission.analysis_mode,
                tool=submission.tool,
                tool_instance=submission.tool_instance,
                type=submission.type,
                event_time=submission.event_time,
                details=submission.details,
                observables=submission.observables,
                tags=submission.tags,
                queue=submission.queue,
                instructions=submission.instructions,
                files=[])

hunt_parser = subparsers.add_parser('hunt')
hunt_sp = hunt_parser.add_subparsers(dest='hunt_cmd')

execute_hunt_parser = hunt_sp.add_parser('execute',
    help="Execute a hunt with the given parameters.")
execute_hunt_parser.add_argument('hunt',
    help="The name of the hunt to execute in the format type:name where type is the hunt type.")
execute_hunt_parser.add_argument('-s', '--start-time', required=False, default=None,
    help="Optional start time. Time spec absolute format is MM/DD/YYYY:HH:MM:SS")
execute_hunt_parser.add_argument('-e', '--end-time', required=False, default=None,
    help="Optional end time. Time spec absolute format is MM/DD/YYYY:HH:MM:SS")
execute_hunt_parser.add_argument('-z', '--timezone', required=False, default=None,
    help="Optional time zone for start time and end time. Defaults to local time zone.")
execute_hunt_parser.add_argument('-v', '--events', required=False, default=False, action='store_true',
    help="Output the events instead of the submissions.")
execute_hunt_parser.add_argument('--json-dir', required=False, default=None,
    help="Store the events as JSON files in the given directory, one per submission created.")
execute_hunt_parser.add_argument('-d', '--details', required=False, default=False, action='store_true',
    help="Include the details of the submissions in the output.")
execute_hunt_parser.add_argument('--submit-alerts', required=False, default=None, 
    help="Submit as alerts to the given host[:port]")
execute_hunt_parser.add_argument('--query-result-file', required=False, default=None,
    help="Valid only for query hunts. Save the raw query results to the given file.")
execute_hunt_parser.set_defaults(func=execute_hunt)

def verify_hunt(args):
    from saq.collectors.hunter import HunterCollector
    from saq.collectors.query_hunter import QueryHunt
    collector = HunterCollector()
    collector.load_hunt_managers()
    failed = False
    for hunt_type, manager in collector.hunt_managers.items():
        manager.load_hunts_from_config()
        if manager.failed_ini_files:
            sys.stderr.write(f"ERROR: unable to load {len(manager.failed_ini_files)} {hunt_type} hunts\n")
            failed = True

    if failed:
        sys.exit(1)

    print("hunt syntax verified")
    sys.exit(0)

verify_hunt_parser = hunt_sp.add_parser('verify',
    help="Verifies that all configured hunts are able to load.")
verify_hunt_parser.set_defaults(func=verify_hunt)

def list_hunts(args):
    from saq.collectors.hunter import HunterCollector
    from saq.collectors.query_hunter import QueryHunt
    collector = HunterCollector()
    collector.load_hunt_managers()
    for hunt_type, manager in sorted(collector.hunt_managers.items()):
        manager.load_hunts_from_config()
        for hunt in sorted(sorted(manager.hunts, key=operator.attrgetter('name')), 
                key=operator.attrgetter('enabled'), reverse=True):
            ini_file = os.path.splitext(os.path.basename(hunt.ini_path))[0]
            status = "E" if hunt.enabled else "D"
            print(f"{status} {hunt_type}:{ini_file} - {hunt.name}")

    sys.exit(0)

list_hunts_parser = hunt_sp.add_parser('list',
    help="""List the available hunts.
    The format of the output is
    E|D type:name - description
    E: enabled
    D: disabled""")
list_hunts_parser.set_defaults(func=list_hunts)

def list_hunt_types(args):
    from saq.collectors.hunter import HunterCollector
    from saq.collectors.query_hunter import QueryHunt
    collector = HunterCollector()
    collector.load_hunt_managers()
    for hunt_type in collector.hunt_managers.keys():
        print(hunt_type)

    sys.exit(0)

list_hunt_types_parser = hunt_sp.add_parser('list-types',
    help="List the available hunting types.")
list_hunt_types_parser.set_defaults(func=list_hunt_types)


# ============================================================================
# persistence data
#

persistence_parser = subparsers.add_parser('persistence', aliases=['per'])
persistence_sp = persistence_parser.add_subparsers(dest='persistence_cmd')

def list_persistence(args):
    from saq.database import Persistence, PersistenceSource
    source_query = saq.db.query(PersistenceSource)
    if args.source:
        source_query = source_query.filter(PersistenceSource.name.like('%{}%'.format(args.source)))

    for source in source_query.order_by(PersistenceSource.name):
        if args.keys or args.name:
            query = saq.db.query(Persistence).filter(Persistence.source_id == source.id)
            if args.name:
                 query = query.filter(Persistence.uuid.like(f'%{args.name}%'))
            if not args.temporal:
                query = query.filter(Persistence.permanent == 1)

            for persistence in query.order_by(Persistence.uuid):
                print(f"{source.name} {persistence.uuid}")
        else:
            print(source.name)

    sys.exit(0)

list_persistence_parser = persistence_sp.add_parser('list',
    help="List all persistence sources.")
list_persistence_parser.add_argument('-k', '--keys', action='store_true', default=False,
    help="Also display permanent keys. Use --search to narrow down results.")
list_persistence_parser.add_argument('-s', '--source',
    help="Search for keys from the given source.")
list_persistence_parser.add_argument('-n', '--name',
    help="Search for key names matching the given pattern.")
list_persistence_parser.add_argument('-t', '--temporal', action='store_true', default=False,
    help="Also display temporal (non-permanent) keys.")
list_persistence_parser.set_defaults(func=list_persistence)

def clear_persistence(args):
    import dateparser
    from saq.database import Persistence, PersistenceSource, and_

    source = saq.db.query(PersistenceSource).filter(PersistenceSource.name == args.source).first()
    if source is None:
        logging.error(f"unknown persistence source {args.source}")
        sys.exit(1)

    if args.all:
        result = saq.db.execute(Persistence.__table__.delete().where(Persistence.source_id == source.id))
        logging.info(f"persistence group {args.source} cleared {result.rowcount} items")

        if not args.dry_run:
            saq.db.commit()

        sys.exit(0)

    elif args.older_than:
        target_date = dateparser.parse(args.older_than)
        result = saq.db.execute(Persistence.__table__.delete().where(and_(Persistence.source_id == source.id,
                                                                          Persistence.permanent == 0,
                                                                          Persistence.last_update < target_date)))

        logging.info(f"persistence group {args.source} cleared {result.rowcount} items")

        if not args.dry_run:
            saq.db.commit()

        sys.exit(0)
        
    for key in args.keys:
        result = saq.db.execute(Persistence.__table__.delete().where(and_(Persistence.source_id == source.id,
                                                                          Persistence.uuid == key)))

        logging.info(f"persistence group {args.source} key {key} cleared {result.rowcount} items")

    if not args.dry_run:
        saq.db.commit()

    sys.exit(0)

clear_persistence_parser = persistence_sp.add_parser('clear',
    help="Clears persistence data for the given source.")
clear_persistence_parser.add_argument('source',
    help="The name of the source to clear.")
clear_persistence_parser.add_argument('keys', nargs="*",
    help="One or more keys to clear.")
clear_persistence_parser.add_argument('--all', action='store_true', default=False,
    help="Clear all persistence data for this source.")
clear_persistence_parser.add_argument('--older-than',
    help="Clear all non-permanent persistence data that is older than a given time.")
clear_persistence_parser.add_argument('--dry-run', action='store_true', default=False,
    help="Do no commit the changes, only report how many would be cleared out.")
    
clear_persistence_parser.set_defaults(func=clear_persistence)

# ============================================================================
# microsoft api utilities
#

msapi_parser = subparsers.add_parser('msapi',
    help="Interface with Microsoft API integrations.")
msapi_sp = msapi_parser.add_subparsers(dest='msapi_cmd')

graph_parser =  msapi_sp.add_parser('graph',
    help="MS Graph API utilites")
graph_sp = graph_parser.add_subparsers(dest='graph_cmd')

def export_azure_devices(args):
    from saq.graph_api import get_api, execute_and_get_all, execute_request

    api = get_api(args.account_name)
    if not api:
        logging.error(f"no graph api instance loaded for {args.account_name}")
        return False
    
    api.initialize()
    headers = {'ConsistencyLevel': 'eventual'}

    if args.device_count:
        headers['Accept'] = 'text/plain'
        url = api.build_url(f'v1.0/devices/$count')
        results = execute_request(api, url, headers=headers)
        print(f"{results} total devices")
        return

    # get all devices by default
    url = api.build_url('v1.0/devices?$top=500')
    if args.device_name:
        url = api.build_url(f'v1.0/devices?$search="displayName:{args.device_name}"')

    results = execute_and_get_all(api, url, headers=headers)
    if not results:
        return None
    if args.stdout or args.device_name:
        print(json.dumps(results, indent=2))
        return
    with open(args.output_file, 'w') as fp:
        fp.write(json.dumps(results))
    if os.path.exists(args.output_file):
        print(f" + wrote: {args.output_file}")

    return

export_azure_devices_parser = graph_sp.add_parser('get-devices',
    help="Getting and exporting azure devices. WARNING: default is to export all devices to json file.")
export_azure_devices_parser.add_argument(
    '--account-name', default='default', 
    help="The company name of a configured Azure tenant."
)
export_azure_devices_parser.add_argument(
    '-d', '--device_name', default=None,
    help="Query Azure AD for this device name."
)
export_azure_devices_parser.add_argument(
    '-c', '--device_count', action='store_true',
    help="Query Azure AD and return a count of devices."
)
export_azure_devices_parser.add_argument(
    '-o', '--output_file', default='azure_device_export.json',
    help="The name of the file to write results to. Default=azure_device_export.json"
)
export_azure_devices_parser.add_argument(
    '--stdout', action='store_true', default=False,
    help="Print JSON output to console."
)

export_azure_devices_parser.set_defaults(func=export_azure_devices)

def list_graph_resources(args):
    from saq.collectors import graph_collections

    grc = graph_collections.GraphResourceCollector()
    grc.load_resources()
    for r in grc.resources:
        print(r)

    sys.exit(0)

validate_resource_parser = graph_sp.add_parser('list-resources',
    help="Validate graph resource configurations.")

validate_resource_parser.set_defaults(func=list_graph_resources)

def print_collection_accounts(args):
    from saq.collectors import graph_collections

    grc = graph_collections.GraphResourceCollector()
    grc.load_collection_accounts()
    for account, value in grc.collection_accounts.items():
        print(f"\t{account}: {value}")

    sys.exit(0)

print_collection_account_parser = graph_sp.add_parser('print-collection-accounts',
    help="print configured collection accounts")
print_collection_account_parser.set_defaults(func=print_collection_accounts)

def get_resource(args):
    from saq.graph_api import GraphAPI
    from saq.collectors import graph_collections

    grc = graph_collections.GraphResourceCollector()

    grc.load_resources()
    if args.resource not in [r.name for r in grc.resources]:
        logging.error(f"no configured resource by the name {args.resource}")
        sys.exit(1)

    resource = [r for r in grc.resources if r.name == args.resource][0]
    logging.info(f"loaded the {resource.name} resource")

    if not resource.ready:
        print("Supply required arguments:")
        for arg in resource.args['required']:
            value = input(f"Set {arg}: ")
            resource.set_argument(arg, value)
        print("Change optional arugments, if you want:")
        for arg in resource.args['optional']:
            default = resource.args['optional'][arg]
            value = input(f"Set {arg} [{default}]: ") or default
            resource.set_argument(arg, value)

    if not grc.build_graph_api_client_map():
        logging.error("failed to get a graph api client.")
        sys.exit(1)

    graph_api_client = grc.graph_api_clients[resource.graph_account_name]
    graph_api_client.initialize()

    results = grc.execute_resource(graph_api_client, resource)
    # NOTE, not handeling pagination here
    if results:
        print(json.dumps(results, indent=2, sort_keys=True))

    sys.exit(0)

graph_resource_parser = graph_sp.add_parser('get',
    help="get a configured resource")
graph_resource_parser.add_argument('resource', action='store',
    help="the name of a configured resource")
graph_resource_parser.set_defaults(func=get_resource)


# ============================================================================
# configuration settings
#

def enable_integration(args):
    import saq.integration
    saq.integration.enable_integration(args.integration)
    sys.exit(0)

def disable_integration(args):
    import saq.integration
    saq.integration.disable_integration(args.integration)
    sys.exit(0)

def list_integrations(args):
    import saq.integration
    saq.integration.list_integrations()
    sys.exit(0)

integration_parser = subparsers.add_parser('integration')
integration_sp = integration_parser.add_subparsers(dest='integration_cmd')

enable_integration_parser = integration_sp.add_parser('enable',
    help="Enables the given integration.")
enable_integration_parser.add_argument('integration',
    help="The integration to enable. Use ace integration list to get the list of available integrations.")
enable_integration_parser.set_defaults(func=enable_integration)

list_integration_parser = integration_sp.add_parser('list',
    help="List the available integrations and show enabled/disabled status.")
list_integration_parser.set_defaults(func=list_integrations)

disable_integration_parser = integration_sp.add_parser('disable',
    help="Disables the given integration.")
disable_integration_parser.add_argument('integration',
    help="The integration to disable. Use ace integration list to get the list of available integrations.")
disable_integration_parser.set_defaults(func=disable_integration)

# ============================================================================
# service control/start/stop
#

service_parser = subparsers.add_parser('service',
    help="Service management commands. See ace service --help for details.")
service_sp = service_parser.add_subparsers(dest='service_cmd')

def start_service(args):
    import saq.service

    # get the list of the services to start
    if args.all:
        service_names = saq.service.get_all_service_names()
    else:
        service_names = args.services

    # replace any dashes with underscores to make it easier to type
    service_names = [_.replace('-', '_') for _ in service_names]

    # you can only debug one service at a time
    if args.debug and len(service_names) > 1:
        logging.error("you can only debug one service at a time")
        sys.exit(1)

    # if no option is specified then the default is to execute in the background
    if not args.daemon and not args.foreground and not args.debug:
        args.daemon = True

    # check service depedencies and reorder startup
    # first iterate over all the dependencies
    # if any are listed that are not running yet AND are not in our list to start
    # then we bail on startup
    dependencies = set() # of service names
    for service_name in service_names:
        for dep_service_name in saq.service.get_service_dependencies(service_name):
            dependencies.add(dep_service_name)
            if saq.service.get_service_status(dep_service_name) != saq.service.SERVICE_STATUS_RUNNING:
                if dep_service_name not in service_names:
                    if args.no_deps:
                        print(f"NOPE: service {service_name} depends on {dep_service_name} and you're not starting it")
                        sys.exit(1)

    service_names.extend(list(dependencies))
    targets = []
    for service_name in service_names:
        targets.append((service_name, saq.service.get_service_dependencies(service_name)))

    # https://stackoverflow.com/a/11564323
    def topological_sort(source):
        """perform topo sort on elements.

        :arg source: list of ``(name, [list of dependancies])`` pairs
        :returns: list of names, with dependancies listed first
        """
        pending = [(name, set(deps)) for name, deps in source] # copy deps so we can modify set in-place       
        emitted = []        
        while pending:
            next_pending = []
            next_emitted = []
            for entry in pending:
                name, deps = entry
                deps.difference_update(emitted) # remove deps we emitted last pass
                if deps: # still has deps? recheck during next pass
                    next_pending.append(entry) 
                else: # no more deps? time to emit
                    yield name 
                    emitted.append(name) # <-- not required, but helps preserve original ordering
                    next_emitted.append(name) # remember what we emitted for difference_update() in next pass
            if not next_emitted: # all entries have unmet deps, one of two things is wrong...
                raise ValueError("cyclic or missing dependancy detected: %r" % (next_pending,))
            pending = next_pending
            emitted = next_emitted

    # go ahead and install signal handlers to stop the services we're about to start up
    services = [] # of ACEService objects
    # at this point we have the list of services we want to start
    for service_name in list(topological_sort(targets)): # in the order they need to be started
        try:
            if saq.service.get_service_status(service_name) == saq.service.SERVICE_STATUS_RUNNING:
                continue

            service = saq.service.get_service_class(service_name)()
            service.start_service(daemon=args.daemon, debug=args.debug, threaded=args.foreground)
            services.append(service)
        except Exception as e:
            logging.error(f"unable to load service {service_name}: {e}")
            traceback.print_exc()
            sys.exit(1)

    # if we're running in daemon mode then we're done here
    if args.daemon:
        sys.exit(0)

    for service in services:
        service.wait_service()

    sys.exit(0)

start_service_parser = service_sp.add_parser('start',
    help="Start ACE services.")
start_service_parser.add_argument('--debug', action='store_true', default=False,
    help="Starts the service in debug mode.")
start_service_parser.add_argument('-d', '--daemon', '--background', action='store_true', default=False,
    help="Starts the service and then executes in the background, tracking the PID. "
         "This is the default if no other option is specified." )
start_service_parser.add_argument('--foreground', action='store_true', default=False,
    help="Starts the service and then executes in the foreground, tracking the PID.")
start_service_parser.add_argument('--all', action='store_true', default=False,
    help="Starts all enabled services.")
start_service_parser.add_argument('--no-deps', action='store_true', default=False,
    help="Do NOT automatically start service dependencies.")
start_service_parser.add_argument('services', nargs="*",
    help="The names of the services to start. Use ace service list to get a list of names.")
start_service_parser.set_defaults(func=start_service)

def _stop_service(args):
    import saq.service

    # get the list of the services to start
    if args.all:
        service_names = saq.service.get_all_service_names()
    else:
        service_names = args.services

    for service_name in service_names:
        service_status = saq.service.get_service_status(service_name)
        if service_status is None:
            logging.error(f"unknown service {service_name}")
        elif service_status == saq.service.SERVICE_STATUS_RUNNING:
            logging.info(f"stopping service {service_name}")
            saq.service.stop_service(service_name)
        else:
            if not args.all:
                logging.info(f"service {service_name} status already ({service_status})")

def stop_service(args):
    _stop_service(args)
    sys.exit(0)

stop_service_parser = service_sp.add_parser('stop',
    help="Stops the services.")
stop_service_parser.add_argument('--all', action='store_true', default=False,
    help="Stops all running services.")
stop_service_parser.add_argument('services', nargs="*",
    help="The names of the service to stop. Use ace service list to get a list of names.")
stop_service_parser.set_defaults(func=stop_service)

def restart_service(args):
    _stop_service(args)
    start_service(args)

restart_service_parser = service_sp.add_parser('restart',
    help="Restart ACE services.")
restart_service_parser.add_argument('--debug', action='store_true', default=False,
    help="Starts the service in debug mode.")
restart_service_parser.add_argument('-d', '--daemon', '--background', action='store_true', default=False,
    help="Starts the service and then executes in the background, tracking the PID. "
         "This is the default if no other option is specified." )
restart_service_parser.add_argument('--foreground', action='store_true', default=False,
    help="Starts the service and then executes in the foreground, tracking the PID.")
restart_service_parser.add_argument('--all', action='store_true', default=False,
    help="Restarts all enabled services.")
restart_service_parser.add_argument('--no-deps', action='store_true', default=False,
    help="Do NOT automatically start service dependencies.")
restart_service_parser.add_argument('services', nargs="*",
    help="The names of the services to restart. Use ace service list to get a list of names.")
restart_service_parser.set_defaults(func=restart_service)

def status_service(args):
    from saq.service import get_all_service_names, get_service_config, get_service_status
    print("{:<30}{:<15}{}".format('SERVICE', 'STATUS', 'DESCRIPTION'))
    for service_name in get_all_service_names():
        service_config = get_service_config(service_name)
        print("{:<30}{:<15}{}".format(service_name, 
                                      get_service_status(service_name),
                                      service_config.get('description', fallback='?')))

status_service_parser = service_sp.add_parser('status',
    help="Lists the available services and their status.")
status_service_parser.set_defaults(func=status_service)

# ============================================================================
# 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
        if f'analysis_mode_{root.analysis_mode}' not in saq.CONFIG:
            logging.error(f"invalid analysis mode {root.analysis_mode}")
            sys.exit(1)

        # disable cleanup in whatever mode we use
        saq.CONFIG[f'analysis_mode_{root.analysis_mode}']['cleanup'] = 'no'

        root.description = args.description if args.description else 'Command Line Correlation'
        root.instructions = args.instructions
        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)

        for tag in args.tags:
            observable.add_tag(tag)

        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.",
    epilog="Example: ace correlate ipv4 8.8.8.8 -G email -E ip_inspector -E gglsbl_service -D carbon_black_process_analysis_v1 -G common")
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', action='append',
    help="Specify a module name (substring match) to disable. This option can be specified multiple times.")
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', action='append', 
    help="Specify module names (substring match) to enable. This option can be specified multiple times.")
correlate_parser.add_argument('-G', '--enable-module-group', required=False, dest='enable_module_group', action='append',
    help="Module groups to enable by name. Specify for each module group.")
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 [+-]0000 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('--add-tag', required=False, dest='tags', action='append', default=[],
    help="Adds the given tag 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('--instructions', required=False, dest='instructions', default=None,
    help="""A free form string value that gives the analyst instructions on what
        this alert is about and/or how to analyze the data contained in the
        alert.""")
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)

# ============================================================================
# Remediation Utilities
#

remediation_parser = subparsers.add_parser('remediation')
remediation_sp = remediation_parser.add_subparsers(dest='remediation_cmd')

def print_remediation_status(id=None, type=None, status=None, key=None, user_id=None, user=None, print_json=False):
    import json
    from saq.database import Remediation, User
    from sqlalchemy import or_

    query = saq.db.query(Remediation).join(User)

    if id:
        query = query.filter(Remediation.id == id)

    if type:
        query = query.filter(Remediation.type == type)

    if status:
        query = query.filter(Remediation.status == status)

    if key:
        query = query.filter(Remediation.key.ilike(f'%{key}%'))

    if user_id:
        query = query.filter(Remediation.user_id == user_id)

    if user:
        query = query.filter(or_(User.username.ilike(f'%{user}%'), 
                                 User.display_name.ilike(f'%{user}%')))

    remediation_data = query.order_by(Remediation.insert_date.desc())
    if not remediation_data.first():
        print("No remediation data found.")
        return None

    if print_json:
        print(json.dumps([r.json for r in remediation_data], default=str))
        return True

    _keys = remediation_data[0].json.keys()
    for r in remediation_data:
        _r_text = " -> "
        _r_data = r.json
        for k in _keys:
            _r_text += f"{k}:{_r_data[k]} | "
        print(_r_text)

def remediation_status(args):
    print_remediation_status(id=args.id,
                             status=args.status,
                             type=args.type,
                             key=args.key,
                             user=args.user,
                             user_id=args.user_id,
                             print_json=args.json)
    sys.exit(0)

remediation_status_parser = remediation_sp.add_parser('status',
    help="Shows the status of remediation requests.")
remediation_status_parser.add_argument('-s', '--status',
    help="Return results matching the given status.")
remediation_status_parser.add_argument('-t', '--type',
    help="Return results matching the given remediation type.")
remediation_status_parser.add_argument('-k', '--key', 
    help="Return results matching the given remediation target (key).")
remediation_status_parser.add_argument('-u', '--user', 
    help="Return results matching the given username or display name.")
remediation_status_parser.add_argument('--user-id', type=int,
    help="Return results matching the given user_id.")
remediation_status_parser.add_argument('-i', '--id', type=int,
    help="Return results matching the given remediation id.")
remediation_status_parser.add_argument('--json', default=False, action='store_true',
    help="Return results as json.")
remediation_status_parser.set_defaults(func=remediation_status)

def execute_remediation(args):
    import saq
    from saq.remediation import RemediationTarget, load_all_remediators, remediate_target

    # load all remediators
    remediators = load_all_remediators()

    # if no targets were passed, make space
    if not hasattr(args, 'targets'):
        args.targets = []

    # targets from a file
    if args.from_file:
        with open(args.from_file, 'r') as fp:
            for line in fp:
                line = line.strip()
                args.targets.append(line)

    # targets from stdin
    if args.from_stdin:
        for line in sys.stdin:
            line = line.strip()
            args.targets.append(line)

    # load all targets
    targets = [
                 RemediationTarget(args.type,
                                   _t,
                                   action=args.remediation_action,
                                   comment=args.comment)
                 for _t in args.targets
                ]

    results = {}
    for target in targets:
        if len(target.history) > 0:
            print(f" + Target '{target.key}' has {len(target.history)} historical remediations.")

        if args.background:
            # Queue for the Remediation Service
            target.queue()
            continue

        # record it, but do it now
        remediate_target(remediators, target)


    if args.wait:
        request_wait_time = saq.CONFIG['service_remediation'].getint('request_wait_time', fallback=10)
        quit_time = time.time() + request_wait_time
        print(f" + waiting up to {request_wait_time} seconds for {len(targets)} targets to reach a completed state.")

        for target in targets:
            # refresh the remediation history
            target.refresh()
            while target.processing and time.time() < quit_time:
                time.sleep(1)
                target.refresh()

    for target in targets:
        target.refresh()
        if args.background:
            print(f" + Requested {target.last_remediation}")
        else:
            print(f" + Executed {target.last_remediation}")

    sys.exit(0)

def execute_email_remediation(args):
    from saq.remediation.mail import create_email_remediation_key
    setattr(args, 'type', 'email')
    setattr(args, 'targets', [create_email_remediation_key(args.message_id, args.email_address)])
    execute_remediation(args)

# most of the arguments are shared between the remove and restore action
def _add_remediation_common_arguments(_parser):
    _parser.add_argument('-b', '--background', action='store_true', default=False,
        help="Executes the action in the background.")
    _parser.add_argument('-w', '--wait', action='store_true', default=False,
        help="Wait until the remediaiton reaches a completed state OR a timeout occurs.")
    _parser.add_argument('-c', '--comment', default="command line request",
        help="Adds the given comment to the remediation as useful note.")
    _parser.add_argument('--from-file', required=False, default=None,
        help="Read the list of targets, one per line, from the given file.")
    _parser.add_argument('--from-stdin', required=False, default=False, action='store_true',
        help="Read the list of targets, one per line, from the standard input.")

def _add_remediation_type_value_arguments(_parser):
    _parser.add_argument('type',
        help="The type of thing you want to target.")
    _parser.add_argument('targets', nargs='*',
        help="The thing or things you want to target.")

remediation_remove_parser = remediation_sp.add_parser('remove',
    help="Executes a remediation action to remove something.")
_add_remediation_common_arguments(remediation_remove_parser)
_add_remediation_type_value_arguments(remediation_remove_parser)
remediation_remove_parser.set_defaults(remediation_action='remove', func=execute_remediation)

remediation_restore_parser = remediation_sp.add_parser('restore',
    help="Executes a remediation action to restore something previously removed.")
_add_remediation_common_arguments(remediation_restore_parser)
_add_remediation_type_value_arguments(remediation_restore_parser)
remediation_restore_parser.set_defaults(remediation_action='restore', func=execute_remediation)

# utility functions to make it easier to use
def _add_email_remediation_target_arguments(_parser):
    _parser.add_argument('message_id', help="The message_id of the email to remove.")
    _parser.add_argument('email_address', help="The email address (mailbox) to remove the email from.")

remediation_remove_email_parser = remediation_sp.add_parser('remove-email',
    help="Executes a remediation action to remove an email.")
_add_remediation_common_arguments(remediation_remove_email_parser)
_add_email_remediation_target_arguments(remediation_remove_email_parser)
remediation_remove_email_parser.set_defaults(remediation_action='remove', func=execute_email_remediation)

remediation_restore_email_parser = remediation_sp.add_parser('restore-email',
    help="Executes a remediation action to restore an email.")
_add_remediation_common_arguments(remediation_restore_email_parser)
_add_email_remediation_target_arguments(remediation_restore_email_parser)
remediation_restore_email_parser.set_defaults(remediation_action='restore', func=execute_email_remediation)


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

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

    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)


# ============================================================================
# Section for Metrics / Analytics / Statistics
#

try:
    from ace_metrics.cli import build_metric_parser, execute_metric_arguments

    metrics_parser = subparsers.add_parser("metrics", help="ace ecosystem metrics")
    build_metric_parser(metrics_parser)

    def metrics(args):
        """Interface with the ace_metrics lib."""

        from saq.database import get_db_connection

        with get_db_connection() as db:
            return execute_metric_arguments(db, args)

    metrics_parser.set_defaults(func=metrics)

except ModuleNotFoundError:
    pass


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

def start_gui(args):
    import saq
    from app import create_app, db
    from werkzeug.serving import run_simple
    from werkzeug.wsgi import DispatcherMiddleware
    from flask import Flask

    app = create_app()
    app.config['DEBUG'] = True
    app.config['APPLICATION_ROOT'] = '/ace'
    app.wsgi_app = DispatcherMiddleware(app.wsgi_app, {
        app.config['APPLICATION_ROOT']: 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)

    run_simple(saq.CONFIG.get('gui', 'listen_address'), saq.CONFIG.getint('gui', 'listen_port'), app,
               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 aceapi 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'
    app.wsgi_app = DispatcherMiddleware(app.wsgi_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'), app,
               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)

# ============================================================================
# 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
    u.display_name = args.display_name
    u.queue = args.queue

    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)

    if args.password:
        password = args.password
    else:
        password = getpass("Enter password for {}: ".format(u.username))
        confirm = getpass("Confirm password for {}: ".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, display_name, queue ) VALUES ( %s, %s, %s, %s, %s, %s )""", (
            u.username, u.email, u.password_hash, u.timezone, u.display_name, u.queue ))
        db.commit()

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

user_parser = subparsers.add_parser('user',
    help="User management commands.")
user_sp = user_parser.add_subparsers(dest='user_cmd')

add_user_parser = user_sp.add_parser('add',
    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.add_argument('-d', '--display-name', required=False, default=None,
    help="The (optional) display name for the user.")
add_user_parser.add_argument('--password', required=False, default=None,
    help="""Provide the password for the user on the command line. 
    Don't do this unless it's a part of automation.""")
add_user_parser.add_argument('-q', '--queue', required=False, default='default', 
    help="Set the default queue the user is assigned to.")
add_user_parser.set_defaults(func=add_user)

# XXX DEPRECATED
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.add_argument('-d', '--display-name', required=False, default=None,
    help="The (optional) display name for the user.")
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 ))

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

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

        if args.enable:
            c.execute("""UPDATE users SET enabled = True WHERE id = %s""", ( user_id ))

        if args.disable:
            c.execute("""UPDATE users SET enabled = False WHERE id = %s""", ( user_id ))

        db.commit()

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

modify_user_parser = user_sp.add_parser('modify',
    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.add_argument('-q', '--default_queue', required=False, default=None, help="Change the default queue the user is assigned to.")
modify_user_parser.add_argument('-n', '--display-name', required=False, default=None, help="Change the user's display name.")
modify_user_parser.add_argument('--enable', action='store_true', required=False, default=None, help="Enable the user.")
modify_user_parser.add_argument('--disable', action='store_true', required=False, default=None, help="Disable the user.")
modify_user_parser.set_defaults(func=modify_user)

# XXX DEPRECATED
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)

delete_user_parser = user_sp.add_parser('delete',
    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)

def list_users(args):
    import saq
    from saq.database import User
    print("{:<6}{:<15}{:<25}{:<40}{:<25}{}".format('Id', 'User', 'Name', 'Queue', 'Enabled', 'TZ'))
    for user in saq.db.query(User).order_by(User.username):
        print("{:<6}{:<15}{:<25}{:<40}{:<25}{}".format(
              user.id, 
              user.username, 
              '' if user.display_name is None else user.display_name, 
              user.queue,
              'True' if user.enabled else 'False',
              user.timezone))

    sys.exit(0)

list_users_parser = user_sp.add_parser('list',
    help="List existing users in the system.")
list_users_parser.set_defaults(func=list_users)

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

test_parser = subparsers.add_parser('test',
    help="Test commands.")
test_sp = test_parser.add_subparsers(dest='test_cmd')

def test_proxy(args):
    try:
        import requests
        from saq import proxy
        requests.get(args.url, proxies=proxy.proxies(), verify=saq.CONFIG['proxy']['verify'] if 'verify' in saq.CONFIG['proxy'] else False)
        sys.exit(0)
    except Exception as e:
        traceback.print_exc()
        sys.exit(1)

test_proxy_parser = test_sp.add_parser('proxy',
    help="Text proxy access")
test_proxy_parser.add_argument('url',
    help="Test the proxy by accessing the given URL. Any content downloaded is discarded.")
test_proxy_parser.set_defaults(func=test_proxy)

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 'hostname' not in saq.CONFIG[key] or not saq.CONFIG[key]['hostname'] \
            or 'database' not in saq.CONFIG[key] or not saq.CONFIG[key]['database'] \
            or 'username' not in saq.CONFIG[key] or not saq.CONFIG[key]['username'] \
            or 'password' not in saq.CONFIG[key] 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(1)

    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 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()

alert_parser = subparsers.add_parser('alert',
    help="Alert management commands.")
alert_sp = alert_parser.add_subparsers(dest='alert_cmd')

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)

create_alert_parser = alert_sp.add_parser('create', aliases=['new'],
    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 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)

rebuild_index_parser = alert_sp.add_parser('rebuild',
    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)

import_alert_parser = alert_sp.add_parser('import',
    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)

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)

delete_alert_parser = alert_sp.add_parser('delete',
    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)

reset_alert_parser = alert_sp.add_parser('reset',
    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()

# 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)

reset_alert_parser = alert_sp.add_parser('archive',
    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)

add_observable_parser = alert_sp.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)

reload_alert_parser = alert_sp.add_parser('analyze', aliases=['reload'],
    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)

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)

cleanup_alerts_parsers = alert_sp.add_parser('cleanup',
    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 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)

display_alert_parser = alert_sp.add_parser('display',
    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)


def print_file_contents(uuid, submission):
    for file in submission.files:
        file_path = os.path.join(saq.SAQ_HOME, saq.CONFIG['global']['data_dir'], saq.CONFIG['global']['node'], uuid[0:3],
                                 uuid, file)
        if not os.path.exists(file_path):
            logging.warning("storage directory {0} does not exist".format(file_path))
        with open(file_path, 'rb') as f:
            print(repr(f.read()))


def display_submission(args):
    import saq
    from saq.analysis import _JSONEncoder, RootAnalysis
    from saq.submission import Submission, get_submission_target_buffer

    root = RootAnalysis(storage_dir=args.storage_dir)
    root.load()

    submission = Submission(**root.submission)

    if args.submission_cmd == 'submission':
        print(get_submission_target_buffer(submission).decode('utf8'))
    elif args.submission_cmd == 'observable':
        observables_json = json.dumps(submission.observables, indent=True, sort_keys=True, cls=_JSONEncoder)
        print(f'{observables_json}'.encode('utf8', errors='backslashreplace').decode('utf8'))
    elif args.submission_cmd == 'file':
        print_file_contents(args.uuid, submission)
    elif args.submission_cmd == 'all':
        print(get_submission_target_buffer(submission).decode('utf8'))
        print_file_contents(args.uuid, submission)

    sys.exit(0)

# submission parser
submission_parser = subparsers.add_parser('submission',
    help="Submission management commands.")

# submission subparser
submission_sp = submission_parser.add_subparsers(dest='submission_cmd')

# display submission parser
display_submission_parser = submission_sp.add_parser('display',
    help="Displays the contents for all components (submission, observables, files) within a given submission as bytes objects.")
# display submission subparsers
display_submission_sp = display_submission_parser.add_subparsers(dest='submission_cmd')

display_submission_all_parser = display_submission_sp.add_parser('all',
    help="Displays the contents for all components (submission, observables, files) within a given submission as bytes objects.")
display_submission_all_parser.add_argument('storage_dir',
    help="The storage directory of the submission to display")
display_submission_all_parser.set_defaults(func=display_submission)

display_submission_submission_parser = display_submission_sp.add_parser('submission',
    help='Displays the buffer used for scanning a given file submission as a bytes object.')
display_submission_submission_parser.add_argument('storage_dir',
    help="The storage directory of the submission to display")
display_submission_submission_parser.set_defaults(func=display_submission)

display_submission_observable_parser = display_submission_sp.add_parser('observable',
    help="Displays the buffer used for scanning a given observable submission as a bytes object.")
display_submission_observable_parser.add_argument('storage_dir',
    help="The storage directory of the submission to display")
display_submission_observable_parser.set_defaults(func=display_submission)

display_submission_file_parser = display_submission_sp.add_parser('file',
    help="Displays the buffer (raw content) used for scanning a given file submission.")
display_submission_file_parser.add_argument('storage_dir',
    help="The storage directory of the submission to display")
display_submission_file_parser.set_defaults(func=display_submission)


# ============================================================================
# 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())

    with get_db_connection(name="email_archive") as db:
        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)

    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 execute_remediation, execute_restoration, REMEDIATION_TYPE_EMAIL
    from saq.remediation.mail import get_remediation_targets, create_email_remediation_key

    # get a user-id to execute remediation
    user = saq.db.query(User).filter(User.username == args.user_name).first()
    if user is None:
        sys.stderr.write(f"invalid user {args.user_name}\n")
        sys.exit(1)

    user_id = user.id

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

    tuples = []
    if args.message_id_only:
        # attempt to build targets from message-ids only 
        tuples = get_remediation_targets(args.targets)
    else:
        # turn targets into array of tuples
        targets = args.targets
        skip_every_other = range(0, len(targets), 2)
        tuples = [(targets[i], targets[i+1]) for i in skip_every_other]
    print(tuples)
    for message_id, recipient in tuples:
        print(message_id, recipient)
        # generate a key from parameters
        key = create_email_remediation_key(message_id, recipient)
        # remediate or restore?
        target_func = execute_remediation
        if args.restore:
            target_func = execute_restoration
        # execute the remediation
        result = target_func(REMEDIATION_TYPE_EMAIL, key, user_id, saq.COMPANY_ID, comment=args.comment)
        print(f"got result {result.id}")

    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('-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, default='ace',
    help="The username to execute the remediation as. Defaults to the automation (ace) user.")
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 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)

# ============================================================================
# Carbon Black Interface
#

def carbon_black_cloud_utils(args):
    from cbinterface.psc.device import yield_devices
    from saq.carbon_black import CBC_API

    if args.export_active_devices:
        output_file = args.output_file
        devices = [d for d in yield_devices(CBC_API, query="sensorStates:ACTIVE")]
        with open(output_file, 'w') as fp:
            fp.write(json.dumps(devices))
        if os.path.exists(output_file):
            print(f" + wrote: {output_file}")

cbc_parser = subparsers.add_parser('cbc-tools',
    help="Interface with the configured CBC environment")
cbc_parser.add_argument('--export-active-devices', action='store_true', help="export active CBC devices to 'active_carbon_black_cloud_devices.json'")
cbc_parser.add_argument(
    '-o', '--output_file', action='store', default='active_carbon_black_cloud_devices.json',
    help="file name to write results to. default is active_carbon_black_cloud_devices.json."
)

cbc_parser.set_defaults(func=carbon_black_cloud_utils)

# ============================================================================
# asset-tracker utilities
#

from saq.asset_tracker import build_parser as at_parser_builder

def asset_tracker(args):
    from saq.asset_tracker import execute_arguments
    return execute_arguments(args)

asset_tracking_parser = subparsers.add_parser('asset-tracking',
    help="Interface with the Asset tracking database")
at_parser_builder(asset_tracking_parser)

asset_tracking_parser.set_defaults(func=asset_tracker)

# ============================================================================
# ip-inspector utilities
#

def update_maxmind_databases(args):
    from saq import proxy
    from ip_inspector import maxmind

    config = saq.CONFIG['analysis_module_ip_inspector']
    proxies = proxy.proxies() if 'use_proxy' in config and config.getboolean('use_proxy') else None
    license_key = config['license_key'] if 'license_key' in config else args.license_key
    if license_key:
        if not maxmind.update_databases(license_key=license_key, proxies=proxies):
            logging.error("Problem updating MaxMind GeoLite2 databases.")
            sys.exit(1)

    sys.exit(0)

update_maxmind_parser = subparsers.add_parser('update-maxmind-databases',
    help="Update the local MaxMind GeoLite2 databases.")
update_maxmind_parser.add_argument('--license-key', action='store', default=None,
    help="Manually supply a license key.")
update_maxmind_parser.set_defaults(func=update_maxmind_databases)

from ip_inspector.cli import build_parser

ip_inspector_parser = subparsers.add_parser("ip-inspector", help="ACE's IP Inspector environment.")
build_parser(ip_inspector_parser)

def ipinspector(args):
    """Interface with the ip_inspector lib."""
    from saq import proxy
    from ip_inspector.cli import execute

    config = saq.CONFIG['analysis_module_ip_inspector']
    use_proxy = config.getboolean('use_proxy')
    if use_proxy:
        proxies = proxy.proxies()
        os.environ['https_proxy'] = proxies['https']
        os.environ['http_proxy'] = proxies['http']

    return execute(args)

ip_inspector_parser.set_defaults(func=ipinspector)


# TODO move this into integration
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 update_tip_cache(args):
    from saq.tip import tip_factory
    tip = tip_factory()
    tip._build_cache()
    sys.exit(0)

update_tip_cache_parser = subparsers.add_parser('update-tip-cache',
    help="Updates the local TIP cache used by the analysis modules.")
update_tip_cache_parser.set_defaults(func=update_tip_cache)

def export_sip_yara_rules(args):
    import glob
    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, rule_count=0):
        if rule_count > 0:
            return os.path.join(args.dir, "SIP_{}_{}.yar".format(indicator_type, rule_count))
        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)

    yara_files_written_or_known = []

    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))
        
        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" ] }


        # handle writing multiple files with single rule limit is hit
        url = f'/api/indicators?status=Analyzed&bulk=True&type={indicator_type}'
        if sources:
            url += '&sources={}'.format(','.join(sources))

        def convert_sip_indicators_to_yara_strings(indicators, rule_count=0):
            """return indicators that still need to be processed because the limit was reached"""

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

            if rule_count > 0:
                template = template.replace('TEMPLATE_RULE_NAME', 'SIP_{}_{}'.format(re.sub(r'[^a-zA-Z0-9_]', '', indicator_type), rule_count))
            else:
                template = template.replace('TEMPLATE_RULE_NAME', 'SIP_{}'.format(re.sub(r'[^a-zA-Z0-9_]', '', indicator_type)))

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

            # https://github.com/VirusTotal/yara/ -> libyara/include/yara/libyara.h#L83
            stack_size = 16384
            # in reality, this stack size equates to 16382 strings per rule.
            # We will use a nice nice round number.
            single_rule_string_limit = stack_size - 384

            remaining_indicators = []
            
            #def convert_indicators_to_yara_strings(indicators, count=0):
            for indicator in indicators:
                string_buffer = io.StringIO()
                item_id = indicator['id']
                item_value = indicator['value']

                if count == single_rule_string_limit:
                    logging.warning(f"reached default yara single rule stack size limit with indicator id: {item_id}")
                    remaining_indicators = indicators[count:]
                    logging.warning(f"breaking and will write remaining {len(remaining_indicators)} indicators to separate rule.")
                    break

                # 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))
                return 

            # 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)
                return

            output_file = get_yara_filename(indicator_type, rule_count)
            yara_files_written_or_known.append(output_file)

            # make sure something changed
            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))
                    return

            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)
                    yara_files_written_or_known.remove(output_file)
                except Exception as e:
                    logging.error("unable to remove {}: {}".format(output_file, e))

            return remaining_indicators

        indicators_for_processing = sip_client.get(url)
        rule_count=0
        while isinstance(indicators_for_processing, list) and len(indicators_for_processing) > 0:
            indicators_for_processing = convert_sip_indicators_to_yara_strings(indicators_for_processing, rule_count=rule_count)
            rule_count += 1

    # use glob to search for old files that need to be deleted
    # like if a mass indicator turn off happened and fell back on our rule count..
    # if we don't do this, a file would be left out there with duplicate strings and waste compute resources.
    for file_path in glob.glob(f"{args.dir}/*.yar"):
        if file_path not in yara_files_written_or_known:
            logging.warning(f"detected indicator reduction left legacy yar file. deleting '{file_path}")
            os.remove(file_path)

    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 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)

# TODO move to integration
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 comprehensive_cloudphish_cache_clearing(args):
    """Clear cloudphish cached on non-TP URLs.

    If a URL hasn't been in a TP alert in X days, clear it even
    if the URL as an Alert status.

    This scenario happens a lot: A new indicator gets turned on that
    causes a URL to get an Alert status with the cloudphish cache. The
    indicator can get turned off but we will continue to alert on the URL
    until the cache is cleared. The longer the FP URL continues without 
    being cleared the more problems can be introduced if it shows up again.

    IF there really are URLs that remain TP for extended periods of time,
    they should be turned into indicators (of various forms) instead of holding
    onto cached analysis results.
    """
    import ace_api
    from saq.database import get_db_connection

    ace_api.set_default_remote_host(saq.CONFIG['global']['node'])
    ace_api.set_default_ssl_ca_path(saq.CONFIG['SSL']['ca_chain_path'])

    clear_counter = 0

    logging.info(f"getting ALERT status cloudphish cache results\
 inserted before {args.cleanup_days} days ago and looked up with the last {args.lookup_days} days...")

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

        c.execute("""
SELECT
    LOWER(HEX(cloudphish_analysis_results.sha256_url)),
    LOWER(HEX(cloudphish_analysis_results.sha256_content)),
    cloudphish_analysis_results.insert_date,
    cloudphish_url_lookup.url
FROM 
    cloudphish_url_lookup JOIN cloudphish_analysis_results ON cloudphish_url_lookup.sha256_url = cloudphish_analysis_results.sha256_url 
WHERE
    cloudphish_analysis_results.result = 'ALERT'
    AND status = 'ANALYZED'
    AND cloudphish_url_lookup.last_lookup > NOW() - INTERVAL %s DAY
    AND cloudphish_analysis_results.insert_date < NOW() - INTERVAL %s DAY
""", ( args.lookup_days, args.cleanup_days, ))

        results = c.fetchall()

        logging.info(f"processing {len(results)} cloudphish ALERT status cache results for URLs with only FP alerts.")

        url_tp_enumeration_tracking = []
        url_tp_enumeration_tracking_path = os.path.join(saq.SAQ_HOME, saq.CONFIG['global']['data_dir'], 'var/cloudphish_enumerated_truepositives.txt') 
        if os.path.exists(url_tp_enumeration_tracking_path):
            with open(url_tp_enumeration_tracking_path, 'r') as fp:
                url_tp_enumeration_tracking = fp.read().splitlines()

        for sha256_url, sha256_content, insert_date, url in results:
            if url in url_tp_enumeration_tracking:
                logging.info(f"skipping previously evaluated url with true positives: {url}")
                continue
            logging.info(f"enumerating dispositions for {url} ...")
            c.execute("""SELECT DISTINCT a.disposition, COUNT(a.disposition) AS Count FROM
                         observables o JOIN observable_mapping om ON o.id = om.observable_id
                           JOIN alerts a ON a.id = om.alert_id
                         WHERE
                           o.type = 'url' AND o.value = %s AND a.alert_type != 'faqueue'
                         AND a.disposition not in ('FALSE_POSITIVE') GROUP BY a.disposition""", (url, ))
            url_alert_results = c.fetchall()

            if not url_alert_results:
                logging.info(f"URL has only FP alert associations, clearing: {url}")
                if not args.dry_run:
                    try:
                        if ace_api.cloudphish_clear_alert(url=url):
                            logging.info(f"cleared {url}.")
                            clear_counter += 1
                        else:
                            logging.warning(f"unexpected result clearing {url}.")
                    except Exception as e:
                        logging.error(f"problem clearing {url} : {e}")
            else:
                url_tp_enumeration_tracking.append(url)
                for dispo, count in url_alert_results:
                    logging.info(f"\t -> has {count} w/ {dispo}")

    logging.info(f"cleared a total of {clear_counter} url cache results.") 

    with open(url_tp_enumeration_tracking_path, 'w') as fp:
        for url in url_tp_enumeration_tracking:
            fp.write(f"{url}\n")

    if os.path.exists(url_tp_enumeration_tracking_path):
        logging.info(f"wrote TP url enumerations to {url_tp_enumeration_tracking_path}")

clear_cloudphish_cache_parser = subparsers.add_parser('cloudphish-clear-stale-alert-cache',
    help="Clear alert status on alertable but stale/false positive URLs.")
clear_cloudphish_cache_parser.add_argument('--dry-run', required=False, dest='dry_run',
    default=False, action='store_true', help="Just return the results that would be purged")
clear_cloudphish_cache_parser.add_argument('--cleanup-days', type=int, default=14,
    help="Select URLs that are older than X days (defaults to 14 days.)")
clear_cloudphish_cache_parser.add_argument('--lookup-days', type=int, default=30,
    help="Select URLs that have been looked up in the last X days (defaults to 30 days")
clear_cloudphish_cache_parser.set_defaults(func=comprehensive_cloudphish_cache_clearing)

# ============================================================================
# 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_proxies(args):
    import requests
    import saq
    proxy_configs = saq.OTHER_PROXIES.copy()
    proxy_configs['GLOBAL'] = saq.CONFIG['proxy']

    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)

# TODO move this to unit testing!
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

    analysis_module_sections = [section for section in saq.CONFIG.sections() if section.startswith('analysis_module_')]
    for section in sorted(analysis_module_sections):

        # is this module disabled globally?
        # modules that are disable globally are not used anywhere
        if not saq.CONFIG.getboolean(section, 'enabled'):
            logging.debug("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)

# ============================================================================
# encryption/decryption utilities
#

encryption_parser = subparsers.add_parser('encryption', aliases=['enc'],
    help="Encryption management commands.")
encryption_sp = encryption_parser.add_subparsers(dest='enc_cmd')

def set_encryption_password(args):
    from saq.service import get_service_status, SERVICE_STATUS_RUNNING
    if saq.service.get_service_status('ecs') == SERVICE_STATUS_RUNNING:
        print("ERROR: the ecs service is currently running")
        print("stop that service before you change the password")
        sys.exit(1)

    from saq.crypto import set_encryption_password, get_aes_key, InvalidPasswordError, encryption_key_set
    while True:
        current_password = None
        if encryption_key_set() and not args.overwrite:
            current_password = getpass.getpass("Enter the CURRENT encryption password:")
            try:
                get_aes_key(current_password)
            except InvalidPasswordError:
                print("ERROR: invalid password")
                print("if you can't remember you can use the --overwrite option")
                print("but then you won't be able to access anything you've already encrypted")
                continue
            
        if args.password is None:
            password = getpass.getpass("enter the new encryption password:")
            password_2 = getpass.getpass("enter the new encryption password again for verification:")
        else:
            password = password_2 = args.password

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

        break

    key = None
    if args.key:
        while True:
            key = getpass.getpass("enter the primary encryption key password:")
            key_2 = getpass.getpass("enter the primary encryption key password again for verification:")

            if key != key_2:
                logging.error("passwords do not match")
                continue

            break

        from Crypto.Hash import SHA256
        h = SHA256.new()
        h.update(key.encode())
        key = h.digest()

    set_encryption_password(password, old_password=current_password, key=key)
    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.add_argument('-o', '--overwrite', default=False, action='store_true',
    help="Overwrites an existing password without prompting.")
set_encryption_password_parser.add_argument('-k', '--key', default=False, action='store_true',
    help="Use the sha256 hash of a string as the primary encryption key. The input is prompted for.")
set_encryption_password_parser.set_defaults(func=set_encryption_password)

set_encryption_password_parser = encryption_sp.add_parser('set',
    help="Sets the password used to encrypt and decrypt archived emails.")
set_encryption_password_parser.add_argument('-o', '--overwrite', default=False, action='store_true',
    help="Overwrites an existing password without prompting.")
set_encryption_password_parser.add_argument('-k', '--key', default=False, action='store_true',
    help="Use the sha256 hash of a string as the primary encryption key. The input is prompted for.")
set_encryption_password_parser.add_argument('-p', '--password',
    help="Use this for the password instead of prompting. Don't use this option unless you have a good reason to.")
set_encryption_password_parser.set_defaults(func=set_encryption_password)

def test_encryption_password(args):
    import saq.crypto
    if args.password is None:
        password = getpass.getpass("Enter the decryption password:")
    else:
        password = args.password

    try:
        saq.crypto.get_aes_key(password)
        logging.info("password OK")
    except saq.crypto.InvalidPasswordError:
        logging.error("invalid password")
        sys.exit(1)

    sys.exit(0)

test_encryption_password_parser = encryption_sp.add_parser('test',
    help="Tests the given password to see if it matches what is currently set as the encryption password.")
test_encryption_password_parser.add_argument('-p', '--password',
        help="Provide the password on the command line. Only recommended for automation purposes.")
test_encryption_password_parser.set_defaults(func=test_encryption_password)

def encrypt_file(args):
    from saq.crypto import encrypt

    password = None
    if args.password:
        password = args.password
    elif args.prompt:
        password = getpass.getpass("enter encryption password: ")
        verify = getpass.getpass("verify encryption password: ")
        if password != verify:
            print("passwords do not match")
            sys.exit(1)

    encrypt(args.source_path, args.target_path, password=password)
    sys.exit(0)

encrypt_file_parser = encryption_sp.add_parser('encrypt',
    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.add_argument('--password', help="Use the given password to encrypt the file.")
encrypt_file_parser.add_argument('--prompt', action='store_true', default=False, 
    help="Prompt for the password to use to encrypt the file.")
encrypt_file_parser.set_defaults(func=encrypt_file)

def decrypt_file(args):
    from saq.crypto import decrypt

    password = None
    if args.password:
        password = args.password
    elif args.prompt:
        password = getpass.getpass("enter decryption password: ")

    decrypt(args.source_path, args.target_path, password=password)
    sys.exit(0)

decrypt_file_parser = encryption_sp.add_parser('decrypt',
    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.add_argument('--password', help="Use the given password to decrypt the file.")
decrypt_file_parser.add_argument('--prompt', help="Prompt for the password to use to decrypt the file.")
decrypt_file_parser.set_defaults(func=decrypt_file)

config_encryption_parser = encryption_sp.add_parser('config',
    help="Configuration encryption management commands.")
config_encryption_sp = config_encryption_parser.add_subparsers(dest='config_enc_cmd')

def set_encrypted_password(args):
    if saq.ENCRYPTION_PASSWORD is None:
        logging.error("missing encryption password (use the -p option or start the ecs service)\n")
        sys.exit(1)

    while True:
        password = getpass.getpass(f"Enter the data to be encrypted and stored:")
        password_2 = getpass.getpass("Re-enter the value for verification:")

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

        break

    from saq.configuration import encrypt_password
    encrypt_password(args.key, password)
    sys.exit(0)

set_encrypted_password_parser = config_encryption_sp.add_parser('set',
    help="Set a password in the system, storing the value encrypted in the database. You will be prompted for the password.")
set_encrypted_password_parser.add_argument('key',
    help="The key (name) of the password.")
set_encrypted_password_parser.set_defaults(func=set_encrypted_password)

def list_encrypted_passwords(args):
    from saq.configuration import export_encrypted_passwords
    encrypted_passwords = export_encrypted_passwords()
    if not encrypted_passwords:
        print("no passwords have been encrypted")
        sys.exit(0)

    for key, value in encrypted_passwords.items():
        if value is None:
            value = "<encrypted>"

        print(f"{key} = {value}")

    sys.exit(0)

list_encrypted_passwords_parser = subparsers.add_parser('list-encrypted-passwords',
    help="Lists the encrypted passwords.")
list_encrypted_passwords_parser.set_defaults(func=list_encrypted_passwords)

list_encrypted_passwords_parser = config_encryption_sp.add_parser('list',
    help="Lists the encrypted passwords.")
list_encrypted_passwords_parser.set_defaults(func=list_encrypted_passwords)

def delete_encrypted_password(args):
    from saq.configuration import delete_password
    if delete_password(args.key):
        print(f"password deleted")
    else:
        print(f"ERROR: unable to delete password")

    sys.exit(0)

delete_encrypted_password_parser = subparsers.add_parser('delete-encrypted-password',
    help="Deletes a given encrypted password.")
delete_encrypted_password_parser.add_argument('key',
    help="The name of the password to delete.")
delete_encrypted_password_parser.set_defaults(func=delete_encrypted_password)

delete_encrypted_password_parser = config_encryption_sp.add_parser('delete',
    help="Deletes a given encrypted password.")
delete_encrypted_password_parser.add_argument('key',
    help="The name of the password to delete.")
delete_encrypted_password_parser.set_defaults(func=delete_encrypted_password)

def export_encrypted_passwords(args):
    import saq.configuration

    if args.file == '-':
        fp = sys.stdout
    else:
        fp = open(args.file, 'w')

    json.dump(saq.configuration.export_encrypted_passwords(), fp)
    fp.write('\n')
    fp.close()
    sys.exit(0)

export_encrypted_passwords_parser = config_encryption_sp.add_parser('export',
    help="Exports the encrypted passwords to JSON.")
export_encrypted_passwords_parser.add_argument('file',
    help="""The name of the file to store the JSON. 
            Use a file name of - to export to stdout.""")
export_encrypted_passwords_parser.set_defaults(func=export_encrypted_passwords)

def import_encrypted_passwords(args):
    import saq.configuration

    if args.file == '-':
        fp = sys.stdin
    else:
        fp = open(args.file, 'r')

    saq.configuration.import_encrypted_passwords(json.load(fp))
    fp.close()
    sys.exit(0)

import_encrypted_passwords_parser = config_encryption_sp.add_parser('import',
    help="Imports the encrypted passwords JSON generated by the export command.")
import_encrypted_passwords_parser.add_argument('file',
    help="""The name of the JSON export. 
            Use a file name of - to import from stdin.""")
import_encrypted_passwords_parser.set_defaults(func=import_encrypted_passwords)

# ============================================================================
# CrowdStrike API utilities
# TODO support integration with CLI

falcon_parser = subparsers.add_parser('falcon',
    help="Utility commands supporting the Crowdstrike Falcon API.")
falcon_sp = falcon_parser.add_subparsers(dest='falcon_cmd')

falcon_oauth_parser = falcon_sp.add_parser('oauth',
    help="Utilities to manage oauth tokens.")
falcon_oauth_sp = falcon_oauth_parser.add_subparsers(dest='falcon_oauth_cmd')

def falcon_oauth_acquire(args):
    from saq.falcon import FalconAPIClient
    with FalconAPIClient() as client:
        print(client.token)

    sys.exit(0)

falcon_oauth_acquire_token_parser = falcon_oauth_sp.add_parser('acquire',
    help="Acquire an oauth token. Does nothing if a token is already acquired and still valid.")
falcon_oauth_acquire_token_parser.set_defaults(func=falcon_oauth_acquire)

def falcon_oauth_revoke(args):
    from saq.falcon import FalconAPIClient
    with FalconAPIClient() as client:
        client.revoke_oauth2_token()

    sys.exit(0)

falcon_oauth_revoke_token_parser = falcon_oauth_sp.add_parser('revoke',
    help="Revoke an oauth token. Does nothing if a token is not already acquired.")
falcon_oauth_revoke_token_parser.set_defaults(func=falcon_oauth_revoke)

def falcon_rtr_list_sessions(args):
    from saq.falcon import FalconAPIClient
    with FalconAPIClient() as client:
        session_ids = client.api_rtr_get_session_ids()
        print(json.dumps(session_ids))

    sys.exit(0)

falcon_rtr_parser = falcon_sp.add_parser('rtr',
    help="Real Time Response")
falcon_rtr_sp = falcon_rtr_parser.add_subparsers(dest='falcon_rtr_cmd')

falcon_rtr_list_sessions_parser = falcon_rtr_sp.add_parser('list-sessions',
    help="")
falcon_rtr_list_sessions_parser.set_defaults(func=falcon_rtr_list_sessions)

def falcon_rtr_open_session(args):
    from saq.falcon import FalconAPIClient
    with FalconAPIClient() as client:
        result = client.api_rtr_open_session(args.device_id)
        print(json.dumps(result))

    sys.exit(0)

falcon_rtr_open_session_parser = falcon_rtr_sp.add_parser('open-session',
    help="")
falcon_rtr_open_session_parser.add_argument('device_id',
    help="The device to open the session on.")
falcon_rtr_open_session_parser.set_defaults(func=falcon_rtr_open_session)

def falcon_rtr_close_session(args):
    from saq.falcon import FalconAPIClient
    with FalconAPIClient() as client:
        result = client.api_rtr_close_session(args.session_id)
        if result:
            print("session closed")

    sys.exit(0)

falcon_rtr_close_session_parser = falcon_rtr_sp.add_parser('close-session',
    help="")
falcon_rtr_close_session_parser.add_argument('session_id',
    help="The id of the session to close.")
falcon_rtr_close_session_parser.set_defaults(func=falcon_rtr_close_session)

def single_falcon_target(func):
    def wrapper(*args, **kwargs):
        from saq.falcon import ( 
            HostNotFoundError, 
            MultipleHostsFoundError
        )
        try:
            return func(*args, **kwargs)
        except HostNotFoundError:
            sys.stderr.write("ERROR: host not found\n")
            sys.exit(1)
        
        except MultipleHostsFoundError as e:
            sys.stderr.write("ERROR: multiple hosts found (see list below)\n")
            for host in e.hosts:
                sys.stderr.write(f"hostname: {host['hostname']} "
                                 f"local_ip: {host['local_ip']} "
                                 f"external_ip: {host['external_ip']} "
                                 f"mac_address: {host['mac_address']} "
                                 f"device_id: {host['device_id']} "
                                 f"last_seen: {host['last_seen']}\n")

    return wrapper

@single_falcon_target
def falcon_rtr_execute_command(args):
    from saq.falcon import ( 
        FalconAPIClient, 
        HostNotFoundError, 
        MultipleHostsFoundError
    )
    
    with FalconAPIClient() as client:
        with client.open_session(
            hostname=args.hostname,
            ipv4=args.ipv4,
            local_ip=args.local_ip,
            external_ip=args.external_ip,
            mac_address=args.mac_address,
            device_id=args.device_id) as host_session:

            if args.admin:
                results = host_session.execute_admin_command(args.base_command, args.command_string)
            else:
                results = host_session.execute_command(args.base_command, args.command_string)

            for result in results:
                if result['stdout']:
                    sys.stdout.write(result['stdout'])
                if result['stderr']:
                    sys.stderr.write(result['stdout'])

    sys.exit(0)

def _add_falcon_host_spec_args(parser):
    parser.add_argument('--hostname',
        help="The host to execute the command on (specified by hostname.)")
    parser.add_argument('--ipv4',
        help="The host to execute the command on (specified by local or external ip.)")
    parser.add_argument('--local-ip',
        help="The host to execute the command on (specified by local ip.)")
    parser.add_argument('--external-ip',
        help="The host to execute the command on (specified by external ip.)")
    parser.add_argument('--mac-address',
        help="The host to execute the command on (specified by mac address in the format XX-XX-XX-XX-XX-XX.)")
    parser.add_argument('--device-id',
        help="The host to execute the command on (specified by device id.) This overrides other filters.")

falcon_rtr_execute_command_parser = falcon_rtr_sp.add_parser('execute',
    help="")
_add_falcon_host_spec_args(falcon_rtr_execute_command_parser)
falcon_rtr_execute_command_parser.add_argument('base_command',
    help="The command to execute on the host.")
falcon_rtr_execute_command_parser.add_argument('command_string',
    help="The command to execute on the host.")
falcon_rtr_execute_command_parser.add_argument('--admin', action='store_true', default=False,
    help="Execute the command as an admin. Some commands require this.")
falcon_rtr_execute_command_parser.set_defaults(func=falcon_rtr_execute_command)

@single_falcon_target
def falcon_rtr_get_file(args):
    from saq.falcon import ( 
        FalconAPIClient, 
        HostNotFoundError, 
        MultipleHostsFoundError
    )

    with FalconAPIClient() as client:
        with client.open_session(
            hostname=args.hostname,
            ipv4=args.ipv4,
            local_ip=args.local_ip,
            external_ip=args.external_ip,
            mac_address=args.mac_address,
            device_id=args.device_id) as host_session:

            host_session.get_file(args.remote_path, args.local_path)

    sys.exit(0)

falcon_rtr_get_file_parser = falcon_rtr_sp.add_parser('get-file',
    help="")
_add_falcon_host_spec_args(falcon_rtr_get_file_parser)
falcon_rtr_get_file_parser.add_argument('remote_path',
    help="The path of the file to get (relative to the remote system.)")
falcon_rtr_get_file_parser.add_argument('local_path',
    help="The path to save the file to locally.")
falcon_rtr_get_file_parser.set_defaults(func=falcon_rtr_get_file)

def falcon_device_search(args):
    from saq.falcon import FalconAPIClient
    with FalconAPIClient() as client:
        result = client.api_search_devices(args.filter)
        print(json.dumps(result))

    sys.exit(0)

falcon_device_parser = falcon_sp.add_parser('device',
    help="Devices")
falcon_device_sp = falcon_device_parser.add_subparsers(dest='falcon_device_cmd')

falcon_device_search_parser = falcon_device_sp.add_parser('search',
    help="Searches for devices known to Falcon.")
falcon_device_search_parser.add_argument('filter',
    help="The FQL filter to apply to the search.")
falcon_device_search_parser.set_defaults(func=falcon_device_search)

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)

    # bash tab completion for argparse
    argcomplete.autocomplete(parser)
    # parse the command line arguments
    args = parser.parse_args()

    # initialize saq
    import saq
    saq.initialize(
            args=args,
            relative_dir=args.relative_dir)

    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)