# vim: sw=4:ts=4:et
#
# ACE Collectors
# These objects collect things for remote ACE nodes to analyze.
#

import importlib
from datetime import datetime, timedelta
import time
import hashlib
import io
import json
import logging
import os, os.path
import pickle
import queue
import shutil
import socket
import tempfile
import threading
import uuid

import ace_api

import saq
from saq.constants import *
from saq.database import (
        use_db,
        execute_with_retry,
        get_db_connection,
        ALERT
)

from saq.constants import *
from saq.error import report_exception
from saq.persistence import Persistable
from saq.service import ACEService
import saq.settings
from saq.submission import Submission, SubmissionFilter
from saq.util import create_directory, abs_path

import urllib3.exceptions
import requests.exceptions
from yara_scanner import YaraScanner
import yara
import plyara

# some constants used as return values
WORK_SUBMITTED = 1
NO_WORK_AVAILABLE = 2
NO_NODES_AVAILABLE = 3
NO_WORK_SUBMITTED = 4

# test modes
TEST_MODE_STARTUP = 'startup'
TEST_MODE_SINGLE_SUBMISSION = 'single_submission'

def submission_target_buffer(self, 
                             description,
                             analysis_mode,
                             tool,
                             tool_instance,
                             type,
                             event_time,
                             tags,
                             observables,
                             details):
    """Returns the buffer used for scanning submission details as a bytes object."""
    from saq.analysis import _JSONEncoder

    details_json = json.dumps(details, indent=True, sort_keys=True, cls=_JSONEncoder)
    observables_json = json.dumps(observables, indent=True, sort_keys=True, cls=_JSONEncoder)
    return f"""description = {description}
analysis_mode = {analysis_mode}
tool = {tool}
tool_instance = {tool_instance}
type = {type}
event_time = {event_time}
tags = {','.join(tags)}

{observables_json}

{details_json}
""".encode('utf8', errors='backslashreplace')

class RemoteNode(object):
    def __init__(self, id, name, location, any_mode, last_update, analysis_mode, workload_count, company_id=None):
        from saq.engine import translate_node

        self.id = id
        self.name = name
        self.location = translate_node(location)
        self.any_mode = any_mode
        self.last_update = last_update
        self.analysis_mode = analysis_mode
        self.workload_count = workload_count
        self.company_id = company_id

        # the directory that contains any files that to be transfered along with submissions
        self.incoming_dir = os.path.join(saq.DATA_DIR, saq.CONFIG['collection']['incoming_dir'])

    def __str__(self):
        return "RemoteNode(id={},name={},location={})".format(self.id, self.name, self.location)

    def submit(self, submission):
        assert isinstance(submission, Submission)

        # if we are submitting locally then we can bypass the API layer
        if self.name == saq.SAQ_NODE and not saq.CONFIG['collection'].getboolean('force_api'):
            return self.submit_local(submission)
        else:
            return self.submit_remote(submission)

    def submit_local(self, submission):
        """Attempts to submit the given the local engine node."""
        logging.debug(f"submitting {submission} locally")
        root = submission.create_root_analysis()
        root.save()

        # if we received a submission for correlation mode then we go ahead and add it to the database
        if root.analysis_mode == ANALYSIS_MODE_CORRELATION:
            ALERT(root)

        root.schedule()

        return { 'result': root.uuid }

    def submit_remote(self, submission):
        """Attempts to submit the given remote Submission to this node."""
        # we need to convert the list of files to what is expected by the ace_api.submit function
        logging.debug(f"submitting {submission} remotely")

        _files = []
        for f in submission.files:
            # we can optionally designate a path in the remote storage_dir for files by using a tuple if (src, dest)
            if isinstance(f, tuple):
                src_path, dest_name = f
                _files.append((dest_name, open(os.path.join(self.incoming_dir, submission.uuid, os.path.basename(src_path)), 'rb')))
            else:
                _files.append((os.path.basename(f), open(os.path.join(self.incoming_dir, submission.uuid, os.path.basename(f)), 'rb')))

        try:
            result = ace_api.submit(
                submission.description,
                remote_host=self.location,
                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,
                company_id=self.company_id,
                files=_files)

            try:
                result = result['result']
                logging.info("submit remote {} submission {} uuid {}".format(self.location, submission, result['uuid']))
            except Exception as e:
                logging.warning("submission irregularity for {}: {}".format(submission, e))
        finally:
            # make sure we clean up our file descriptors
            for name, fp in _files:
                try:
                    fp.close()
                except Exception as e:
                    logging.error("unable to close file descriptor for {}: {}".format(name, e))

        return result

class RemoteNodeGroup(object):
    """Represents a collection of one or more RemoteNode objects that share the
       same group configuration property."""

    def __init__(
            self,
            name,
            coverage,
            full_delivery,
            company_id,
            database,
            group_id,
            workload_type_id,
            shutdown_event,
            batch_size=32,
            target_node_as_company_id=None,
            target_nodes=[],
            thread_count=1):

        assert isinstance(name, str) and name
        assert isinstance(coverage, int) and coverage > 0 and coverage <= 100
        assert isinstance(full_delivery, bool)
        assert isinstance(company_id, int)
        assert isinstance(database, str)
        assert isinstance(group_id, int)
        assert isinstance(workload_type_id, int)
        assert isinstance(shutdown_event, threading.Event)
        assert isinstance(target_nodes, list)
        assert isinstance(thread_count, int)

        self.name = name

        # this the percentage of submissions that are actually sent to this node group
        self.coverage = coverage
        self.coverage_counter = 0

        # if full_delivery is True then all submissions assigned to the group will eventually be submitted
        # if set to False then at least one attempt is made to submit
        # setting to False is useful for QA and development type systems
        self.full_delivery = full_delivery

        # the company this node group belongs to
        self.company_id = company_id

        # A company id for the primary node sharing this company data
        self.target_node_as_company_id = target_node_as_company_id

        # the name of the database to query for node status
        self.database = database

        # the id of this group in the work_distribution_groups table
        self.group_id = group_id

        # the type of work that this collector works with
        self.workload_type_id = workload_type_id

        # the (maximum) number of work items to pull at once from the database
        self.batch_size = batch_size

        # metrics
        self.assigned_count = 0 # how many emails were assigned to this group
        self.skipped_count = 0 # how many emails have skipped due to coverage rules
        self.delivery_failures = 0 # how many emails failed to delivery when full_delivery is disabled

        # total number of threads to run for submission
        self.thread_count = thread_count
        # main threads of execution for this group
        self.threads = []

        # reference to Controller.shutdown_event, used to synchronize a clean shutdown
        self.shutdown_event = shutdown_event

        # an optional list of target nodes names this group will limit itself to
        # if this list is empty then there is no limit
        self.target_nodes = target_nodes

        # when do we think a node has gone offline
        # each node (engine) should update it's status every [engine][node_status_update_frequency] seconds
        # so we wait for twice that long until we think a node is offline
        # at which point we no longer consider it for submissions
        self.node_status_update_frequency = saq.CONFIG['service_engine'].getint('node_status_update_frequency')

        # the directory that contains any files that to be transfered along with submissions
        self.incoming_dir = os.path.join(saq.DATA_DIR, saq.CONFIG['collection']['incoming_dir'])
        
        # sync lock for assigning work to the threads
        self.work_sync_lock = threading.RLock()

    def start(self):
        self.shutdown_event.clear()
        self.clear_work_locks()

        # main threads of execution for this group
        for index in range(self.thread_count):
            thread = threading.Thread(target=self.loop, args=(str(uuid.uuid4()),), name=f"RemoteNodeGroup {self.name} - {index}")
            thread.start()
            self.threads.append(thread)

    def stop(self):
        self.shutdown_event.set()

    def wait(self):
        for thread in self.threads:
            logging.debug(f"waiting for {thread} to complete")
            thread.join()

        self.threads = []

    def loop(self, work_lock_uuid):
        while True:
            try:
                result = self.execute(work_lock_uuid)

                # if we did something then we immediately look for more work unless we're shutting down
                if result == WORK_SUBMITTED:
                    if self.shutdown_event.is_set():
                        break
                # if were was no work available to be submitted then wait a second and look again
                elif result == NO_WORK_AVAILABLE:
                    if self.shutdown_event.wait(1):
                        break
                # if there were no NODES available then wait a little while longer and look again
                elif result == NO_NODES_AVAILABLE:
                    if self.shutdown_event.wait(self.node_status_update_frequency / 2):
                        break
                elif result == NO_WORK_SUBMITTED:
                    if self.shutdown_event.wait(1):
                        break

            except Exception as e:
                logging.error("unexpected exception thrown in loop for {}: {}".format(self, e))
                report_exception()
                if self.shutdown_event.wait(1):
                    break

            finally:
                saq.db.remove()

    @use_db(name='collection')
    def execute(self, work_lock_uuid, db, c):
        # first we get a list of all the distinct analysis modes available in the work queue
        c.execute("""
SELECT DISTINCT(incoming_workload.mode)
FROM
    incoming_workload JOIN work_distribution ON incoming_workload.id = work_distribution.work_id
WHERE
    incoming_workload.type_id = %s
    AND work_distribution.group_id = %s
    AND work_distribution.status IN ( 'READY', 'LOCKED' )
""", (self.workload_type_id, self.group_id,))
        available_modes = c.fetchall()
        db.commit()

        # if we get nothing from this query then no work is available for this group
        if not available_modes:
            if saq.UNIT_TESTING:
                logging.debug("no work available for {}".format(self))
            return NO_WORK_AVAILABLE

        # flatten this out to a list of analysis modes
        available_modes = [_[0] for _ in available_modes]

        # given this list of modes that need remote targets, see what is currently available
        with get_db_connection(self.database) as node_db:
            node_c = node_db.cursor()

            sql = """
SELECT
    nodes.id, 
    nodes.name, 
    nodes.location, 
    nodes.any_mode,
    nodes.last_update,
    node_modes.analysis_mode,
    COUNT(workload.id) AS 'WORKLOAD_COUNT'
FROM
    nodes LEFT JOIN node_modes ON nodes.id = node_modes.node_id
    LEFT JOIN node_modes_excluded ON nodes.id = node_modes_excluded.node_id
    LEFT JOIN workload ON nodes.id = workload.node_id
WHERE
    {where_clause}
GROUP BY
    nodes.id,
    nodes.name,
    nodes.location,
    nodes.any_mode,
    nodes.last_update,
    node_modes.analysis_mode,
    node_modes_excluded.analysis_mode
ORDER BY
    WORKLOAD_COUNT ASC,
    nodes.last_update ASC
"""
            where_clause = []
            where_clause_params = []

            # XXX not sure what this does
            company_id = self.company_id
            if self.target_node_as_company_id is not None:
                company_id = self.target_node_as_company_id

            where_clause.append("nodes.company_id = %s")
            where_clause_params.append(company_id)

            where_clause.append("nodes.is_local = 0")

            where_clause.append("TIMESTAMPDIFF(SECOND, nodes.last_update, NOW()) <= %s")
            where_clause_params.append(self.node_status_update_frequency * 2)

            param_str = ','.join(['%s' for _ in available_modes])
            where_clause.append(f""" 
            (
                (nodes.any_mode AND 
                    (node_modes_excluded.analysis_mode IS NULL 
                     OR node_modes_excluded.analysis_mode NOT IN ( {param_str} )
                    )
                )
                OR node_modes.analysis_mode IN ( {param_str} )
            ) """)
            where_clause_params.extend(available_modes)
            where_clause_params.extend(available_modes)

            # are we limiting what nodes we are sending to?
            if self.target_nodes:
                param_str = ','.join(['%s' for _ in self.target_nodes])
                where_clause.append(f"nodes.name IN ( {param_str} )")
                where_clause_params.extend(self.target_nodes)

            sql = sql.format(where_clause='AND '.join([f'( {_} ) ' for _ in where_clause]))
            #logging.debug(f"MARKER: {sql} {where_clause_params}")
            node_c.execute(sql, tuple(where_clause_params))
            node_status = node_c.fetchall()

        if not node_status:
            logging.warning("no remote nodes are avaiable for all analysis modes {} for {}".format(
                            ','.join(available_modes), self))

            if not self.full_delivery:
                # if this node group is NOT in full_delivery mode and there are no nodes available at all
                # then we just clear out the work queue for this group
                # if this isn't done then the work will pile up waiting for a node to come online
                execute_with_retry(db, c, "UPDATE work_distribution SET status = 'ERROR' WHERE group_id = %s",
                                  (self.group_id,), commit=True)

            return NO_NODES_AVAILABLE

        # now figure out what analysis modes are actually available for processing
        analysis_mode_mapping = {} # key = analysis_mode, value = [ RemoteNode ]
        any_mode_nodes = [] # list of nodes with any_mode set to True
        
        for node_id, name, location, any_mode, last_update, analysis_mode, workload_count in node_status:
            remote_node = RemoteNode(node_id, name, location, any_mode, last_update, analysis_mode, workload_count, company_id=self.company_id)
            if any_mode:
                any_mode_nodes.append(remote_node)

            if analysis_mode:
                if analysis_mode not in analysis_mode_mapping:
                    analysis_mode_mapping[analysis_mode] = []

                analysis_mode_mapping[analysis_mode].append(remote_node)

        # now we trim our list of analysis modes down to what is available
        # if we don't have a node that supports any mode
        if not any_mode_nodes:
            available_modes = [m for m in available_modes if m in analysis_mode_mapping.keys()]
            logging.debug("available_modes = {} after checking available nodes".format(available_modes))

        if not available_modes:
            logging.debug("no nodes are available that support the available analysis modes")
            return NO_NODES_AVAILABLE

        # do we have anything locked yet?
        c.execute("SELECT COUNT(*) FROM work_distribution WHERE lock_uuid = %s AND status IN ( 'READY', 'LOCKED' )", (work_lock_uuid,))
        result = c.fetchone()
        lock_count = result[0]

        if lock_count > 0:
            logging.debug(f"already have {lock_count} work items locked by {work_lock_uuid}")

        # if we don't have any locks yet, go make some
        if lock_count == 0:
            sql = """
UPDATE work_distribution
SET
    status = 'LOCKED',
    lock_time = NOW(),
    lock_uuid = %s
WHERE 
    group_id = %s
    AND work_id IN ( SELECT * FROM ( 
        SELECT
            incoming_workload.id
        FROM
            incoming_workload JOIN work_distribution ON incoming_workload.id = work_distribution.work_id
        WHERE
            incoming_workload.type_id = %s
            AND work_distribution.group_id = %s
            AND incoming_workload.mode IN ( {} )
            AND (
                work_distribution.status = 'READY'
                OR ( work_distribution.status = 'LOCKED' AND TIMESTAMPDIFF(minute, work_distribution.lock_time, NOW()) >= 10 )
            )
        ORDER BY
            incoming_workload.id ASC
        LIMIT %s ) AS t1 )
""".format(','.join(['%s' for _ in available_modes]))
            params = [ work_lock_uuid, self.group_id, self.workload_type_id, self.group_id ]
            params.extend(available_modes)
            params.append(self.batch_size)

            #logging.info(f"MARKER: {sql} {params}")
            with self.work_sync_lock:
                execute_with_retry(db, c, sql, tuple(params), commit=True)

        # now we get the next things to submit from the database that have an analysis mode that is currently
        # available to be submitted to

        sql = """
SELECT 
    incoming_workload.id,
    incoming_workload.mode,
    incoming_workload.work
FROM
    incoming_workload JOIN work_distribution ON incoming_workload.id = work_distribution.work_id
WHERE
    work_distribution.lock_uuid = %s AND work_distribution.status = 'LOCKED'
ORDER BY
    incoming_workload.id ASC
"""
        params = [ work_lock_uuid ]
        c.execute(sql, tuple(params))
        work_batch = c.fetchall()
        db.commit()

        if len(work_batch) > 0:
            logging.info("submitting {} items".format(len(work_batch)))

        # simple flag that gets set if ANY submission is successful
        submission_success = False

        # we should have a small list of things to submit to remote nodes for this group
        for work_id, analysis_mode, submission_blob in work_batch:
            logging.info(f"preparing workload {work_id}")
            # first make sure we can un-pickle this
            try:
                submission = pickle.loads(submission_blob)
            except Exception as e:
                execute_with_retry(db, c, """UPDATE work_distribution SET status = 'ERROR' 
                                             WHERE group_id = %s AND work_id = %s""",
                                  (self.group_id, self.work_id), commit=True)
                logging.error("unable to un-pickle submission blob for id {}: {}".format(work_id, e))
                continue

            # simple flag to remember if we failed to send
            submission_failed = False

            # the result of the submission (we pass to Submission.success later)
            submission_result = None
                
            self.coverage_counter += self.coverage
            if self.coverage_counter < 100:
                # we'll be skipping this one
                logging.debug("skipping work id {} for group {} due to coverage constraints".format(
                              work_id, self.name))
            else:
                # otherwise we try to submit it
                self.coverage_counter -= 100

                # sort the list of RemoteNode objects by the workload_count
                available_targets = any_mode_nodes[:]
                if analysis_mode in analysis_mode_mapping:
                    available_targets.extend(analysis_mode_mapping[analysis_mode])
            
                target = sorted(available_targets, key=lambda n: n.workload_count)
                target = target[0] 

                # attempt the send
                try:
                    submission_result = target.submit(submission)
                    logging.info("{} got submission result {} for {}".format(self, submission_result, submission))
                    submission_success = True
                except Exception as e:
                    log_function = logging.error
                    if not self.full_delivery:
                        log_function = logging.warning
                    else:
                        if not isinstance(e, urllib3.exceptions.MaxRetryError) \
                                and not isinstance(e, urllib3.exceptions.NewConnectionError) \
                                and not isinstance(e, requests.exceptions.ConnectionError):
                            # if it's not a connection issue then report it
                            report_exception()

                    log_function("unable to submit work item {} to {} via group {}: {}".format(
                            submission, target, self, e))

                    # if we are in full delivery mode then we need to try this one again later
                    if self.full_delivery and (isinstance(e, urllib3.exceptions.MaxRetryError) \
                                               or isinstance(e, urllib3.exceptions.NewConnectionError) \
                                               or isinstance(e, requests.exceptions.ConnectionError)):
                        continue

                    # otherwise we consider it a failure
                    submission_failed = True
                    execute_with_retry(db, c, """UPDATE work_distribution SET status = 'ERROR' 
                                                 WHERE group_id = %s AND work_id = %s""",
                                       (self.group_id, work_id), commit=True)
            
            # if we skipped it or we sent it, then we're done with it
            if not submission_failed:
                execute_with_retry(db, c, """UPDATE work_distribution SET status = 'COMPLETED' 
                                             WHERE group_id = %s AND work_id = %s""",
                                  (self.group_id, work_id), commit=True)

            if submission_failed:
                try:
                    submission.fail(self)
                except Exception as e:
                    logging.error(f"call to {submission}.fail() failed: {e}")
                    report_exception()
            else:
                try:
                    submission.success(self, submission_result)
                except Exception as e:
                    logging.error(f"call to {submission}.success() failed: {e}")
                    report_exception()

        if submission_success:
            return WORK_SUBMITTED

        return NO_WORK_SUBMITTED

    @use_db(name="collection")
    def clear_work_locks(self, db, c):
        """Clears any work locks set with work assigned to this group."""
        c.execute("""
        UPDATE work_distribution SET 
            status = 'READY', 
            lock_uuid = NULL, 
            lock_time = NULL 
        WHERE 
            status = 'LOCKED' AND group_id = %s
        """, (self.group_id,))
        db.commit()

    def __str__(self):
        return "RemoteNodeGroup(name={}, coverage={}, full_delivery={}, company_id={}, database={})".format(
                self.name, self.coverage, self.full_delivery, self.company_id, self.database)

class Collector(ACEService, Persistable):
    def __init__(self, workload_type=None, 
                       delete_files=False, 
                       test_mode=None, 
                       collection_frequency=1, 
                       *args, **kwargs):

        super().__init__(*args, **kwargs)

        # often used as the "tool_instance" property of analysis
        self.fqdn = socket.getfqdn()

        # the type of work this collector collects
        # this maps to incoming_workload_type.name in the database
        self.workload_type = workload_type

        # the list of RemoteNodeGroup targets this collector will send to
        self.remote_node_groups = []

        # the directory that contains any files that to be transfered along with submissions
        self.incoming_dir = os.path.join(saq.DATA_DIR, saq.CONFIG['collection']['incoming_dir'])

        # the id of the persistence source this collector uses to store persistence data in the database
        # loaded at initialization
        self.persistence_source_id = None

        # the directory that can contain various forms of persistence for collections
        self.persistence_dir = os.path.join(saq.DATA_DIR, saq.CONFIG['collection']['persistence_dir'])

        # if delete_files is True then any files copied for submission are deleted after being
        # successfully added to the submission queue
        # this is useful for collectors which are supposed to consume and clear the input
        self.delete_files = delete_files

        # test_mode gets set during unit testing
        self.test_mode = test_mode
        if self.test_mode is not None:
            logging.info("*** COLLECTOR {} STARTED IN TEST MODE {} ***".format(self, self.test_mode))

        # the (optional) list of Submission items to send to remote nodes
        # NOTE that subclasses can simply override get_next_submission and not use this queue
        self.submission_list = queue.Queue()

        # the total number of submissions sent to the RemoteNode objects (added to the incoming_workload table)
        self.submission_count = 0

        # primary collection thread that pulls Submission objects to be sent to remote nodes
        self.collection_thread = None

        # repeatedly calls execute_workload_cleanup
        self.cleanup_thread = None

        # optional thread a subclass can use (by overriding the extended_collection functions)
        self.extended_collection_thread = None

        # when was the persistent data last cleared
        self.persistent_clear_time = time.time()

        # how often to collect, defaults to 1 second
        # NOTE there is no wait if something was previously collected
        self.collection_frequency = collection_frequency

        # this is used to filter out submissions according to yara rules
        # see README.SUBMISSION_FILTERS
        self.submission_filter = SubmissionFilter()

        # XXX meh -- maybe this should be hard coded, or at least in a configuratin file or something
        # get the workload type_id from the database, or, add it if it does not already exist
        try:
            with get_db_connection('collection') as db:
                c = db.cursor()
                c.execute("SELECT id FROM incoming_workload_type WHERE name = %s", (self.workload_type,))
                row = c.fetchone()
                if row is None:
                    c.execute("INSERT INTO incoming_workload_type ( name ) VALUES ( %s )", (self.workload_type,))
                    db.commit()
                    c.execute("SELECT id FROM incoming_workload_type WHERE name = %s", (self.workload_type,))
                    row = c.fetchone()
                    if row is None:
                        raise ValueError("unable to create workload type for {}".format(self.workload_type))

                self.workload_type_id = row[0]
                logging.debug("got workload type id {} for {}".format(self.workload_type_id, self.workload_type))

        except Exception as e:
            logging.critical("unable to get workload type_id from database: {}".format(self.workload_type))
            raise e

    def initialize_collector(self):
        """Called automatically at the end of initialize_environment."""
        pass

    def queue_submission(self, submission, key=None):
        """Adds the given Submission object to the queue."""
        assert isinstance(submission, Submission)
        submit = True
        if key is not None:
            key_hash = hashlib.md5(key.encode("utf-8")).hexdigest()
            if self.persistent_data_exists(key_hash):
                submit = False
            self.save_persistent_key(key_hash)

        if submit:
            self.prepare_submission_files(submission)
            self.submission_list.put(submission)

    # 
    # ACEService implementation
    # ------------

    def execute_service(self):
        if self.service_is_debug:
            return self.debug()
        else:
            self.start()

        if not self.service_is_debug:
            self.wait()

    def wait_service(self, *args, **kwargs):
        super().wait_service(*args, **kwargs)
        if not self.service_is_debug:
            self.wait()

    def initialize_service_environment(self):

        # make sure these directories exist
        for dir_path in [ self.incoming_dir, self.persistence_dir ]:
            create_directory(dir_path)

        # load the remote node groups if we haven't already
        if not self.remote_node_groups:
            self.load_groups()

        # make sure at least one is loaded
        if not self.remote_node_groups:
            raise RuntimeError("no RemoteNodeGroup objects have been added to {}".format(self))

        # load tuning rules
        self.submission_filter.load_tuning_rules()

        # initialize persistence for this collector
        self.register_persistence_source(self.service_name)

        # call any subclass-defined initialization routines
        self.initialize_collector()

    #
    # ------------

    def start(self):

        # if we're starting and we haven't loaded any groups yet then go ahead and load them here
        if not self.remote_node_groups:
            self.load_groups()

        self.collection_thread = threading.Thread(target=self.loop, name="Collector")
        self.collection_thread.start()

        self.cleanup_thread = threading.Thread(target=self.cleanup_loop, name="Collector Cleanup")
        self.cleanup_thread.start()

        self.extended_collection_thread = threading.Thread(target=self.extended_collection, name="Extended")
        self.extended_collection_thread.start()

        # start the node groups
        for group in self.remote_node_groups:
            group.start()

    def debug(self):
        # if we're starting and we haven't loaded any groups yet then go ahead and load them here
        if not self.remote_node_groups:
            self.load_groups()

        try:
            self.debug_extended_collection()
        except NotImplementedError:
            pass

        self.execute()
        self.execute_workload_cleanup()

        # start the node groups
        #for group in self.remote_node_groups:
            #group.start()

    def stop(self, *args, **kwargs):
        return self.stop_service(*args, **kwargs)

    def wait(self):
        logging.info("waiting for collection thread to terminate...")
        self.collection_thread.join()
        for group in self.remote_node_groups:
            logging.info("waiting for {} thread to terminate...".format(group))
            group.wait()

        logging.info("waiting for cleanup thread to terminate...")
        self.cleanup_thread.join()

        logging.info("waiting for extended collection thread to terminate...")
        self.extended_collection_thread.join()

        logging.info("collection ended")

    def add_group(
        self, 
        name, 
        coverage, 
        full_delivery, 
        company_id, 
        database, 
        batch_size=32,
        target_node_as_company_id=None,
        target_nodes=[],
        thread_count=1):

        with get_db_connection('collection') as db:
            c = db.cursor()
            c.execute("SELECT id FROM work_distribution_groups WHERE name = %s", (name,))
            row = c.fetchone()
            if row is None:
                c.execute("INSERT INTO work_distribution_groups ( name ) VALUES ( %s )", (name,))
                group_id = c.lastrowid
                db.commit()
            else:
                group_id = row[0]

            remote_node_group = RemoteNodeGroup(
                name, 
                coverage, 
                full_delivery, 
                company_id, 
                database, 
                group_id, 
                self.workload_type_id, 
                self.service_shutdown_event, 
                batch_size=batch_size,
                target_node_as_company_id=target_node_as_company_id,
                target_nodes=target_nodes,
                thread_count=thread_count)
            self.remote_node_groups.append(remote_node_group)
            logging.info("added {}".format(remote_node_group))
            return remote_node_group

    def load_groups(self):
        """Loads groups from the ACE configuration file."""
        for section in saq.CONFIG.keys():
            if not section.startswith('collection_group_'):
                continue

            if not saq.CONFIG[section].getboolean('enabled', fallback=True):
                logging.debug(f"collection group {section} disabled")
                continue

            group_name = section[len('collection_group_'):]
            coverage = saq.CONFIG[section].getint('coverage')
            full_delivery = saq.CONFIG[section].getboolean('full_delivery')
            company_id = saq.CONFIG[section].getint('company_id')
            database = saq.CONFIG[section]['database']
            batch_size = saq.CONFIG[section].getint('batch_size', fallback=32)
            thread_count = saq.CONFIG[section].getint('thread_count', fallback=1)
            
            target_node_as_company_id = None
            if 'target_node_as_company_id' in saq.CONFIG[section]:
                target_node_as_company_id = saq.CONFIG[section]['target_node_as_company_id']

            target_nodes = []
            if 'target_nodes' in saq.CONFIG[section]:
                for node in saq.CONFIG[section]['target_nodes'].split(','):
                    if not node:
                        continue

                    if node == 'LOCAL':
                        node = saq.SAQ_NODE

                    target_nodes.append(node)

            logging.info("loaded group {} coverage {} full_delivery {} company_id {} database {} target_node_as_company_id {} target_nodes {} thread_count {} batch_size {}".format(
                         group_name, coverage, full_delivery, company_id, database, target_node_as_company_id, target_nodes, thread_count, batch_size))

            self.add_group(
                    group_name,
                    coverage,
                    full_delivery,
                    company_id,
                    database,
                    batch_size=batch_size,
                    target_node_as_company_id=target_node_as_company_id,
                    target_nodes=target_nodes,
                    thread_count=thread_count)

    def _signal_handler(self, signum, frame):
        self.stop()

    def initialize(self):
        pass

    def cleanup_loop(self):
        logging.debug("starting cleanup loop")

        while True:
            wait_time = 1
            try:
                if self.execute_workload_cleanup() > 0:
                    wait_time = 0

            except Exception as e:
                logging.exception(f"unable to execute workload cleanup: {e}")

            if self.service_shutdown_event.wait(wait_time):
                break

        logging.debug("exited cleanup loop")

    @use_db(name='collection')
    def execute_workload_cleanup(self, db, c):
        # look up all the work that is currently completed
        # a completed work item has no entries in the work_distribution table with a status in ('READY', 'LOCKED')
        c.execute("""
SELECT 
    i.id, 
    i.work
FROM 
    incoming_workload i JOIN work_distribution w ON i.id = w.work_id
    JOIN incoming_workload_type t ON i.type_id = t.id
WHERE
    t.id = %s
GROUP BY 
    i.id, i.work
HAVING
    SUM(IF(w.status IN ('READY', 'LOCKED'), 1, 0)) = 0
LIMIT 100""", (self.workload_type_id,))

        submission_count = 0
        rows = c.fetchall()
        db.commit()

        for work_id, submission_blob in rows:
            submission_count += 1
            submission = None

            try:
                submission = pickle.loads(submission_blob)
            except Exception as e:
                logging.error(f"unable to un-pickle submission blob for id {work_id}: {e}")

            # clear any files that back the submission
            if submission and submission.files:
                try:
                    target_dir = os.path.join(self.incoming_dir, submission.uuid)
                    shutil.rmtree(target_dir)
                    logging.debug(f"deleted incoming dir {target_dir}")
                except Exception as e:
                    logging.error(f"unable to delete directory {target_dir}: {e}")

            # we finally clear the database entry for this workload item
            execute_with_retry(db, c, "DELETE FROM incoming_workload WHERE id = %s", (work_id,), commit=True)
            logging.info(f"completed work item {work_id}")

        return submission_count

    def loop(self):
        while True:
            try:
                self.execute()
            except Exception as e:
                logging.error("unexpected exception thrown during loop for {}: {}".format(self, e))
                report_exception()
                if self.service_shutdown_event.wait(1):
                    break
            finally:
                saq.db.remove()

            if self.is_service_shutdown:
                break

    def execute(self):

        if self.test_mode == TEST_MODE_STARTUP:
            next_submission = None
        elif self.test_mode == TEST_MODE_SINGLE_SUBMISSION and self.submission_count > 0:
            next_submission = None
        else:
            next_submission = self.get_next_submission()

        # did we not get anything to submit?
        if next_submission is None:
            if self.service_is_debug:
                return

            # wait until we check again (defaults to 1 second, passed in on constructor)
            self.service_shutdown_event.wait(self.collection_frequency)
            return

        if not isinstance(next_submission, Submission):
            logging.critical("get_next_submission() must return an object derived from Submission")

        # does this submission match any tuning rules we have?
        tuning_matches = self.submission_filter.get_tuning_matches(next_submission)
        if tuning_matches:
            self.submission_filter.log_tuning_matches(next_submission, tuning_matches)
            self.cleanup_submission(next_submission)
            return

        self.prepare_submission_files(next_submission)
        self.schedule_submission(next_submission)
        self.cleanup_submission(next_submission)

    def prepare_submission_files(self, submission):
        if not submission.files:
            return

        # XXX hack -- refactor out
        if submission.files_prepared:
            return

        try:
            if not os.path.exists(submission.storage_dir):
                create_directory(submission.storage_dir)

            # move or copy the files into the storage directory of the submission
            # then update the file list with the new paths
            updated_files = []
            for file_submission in submission.files:
                # this could be a tuple of (source_file, target_name)
                if isinstance(file_submission, tuple):
                    source_path, dest_path = file_submission
                else:
                    source_path = file_submission

                target_path = os.path.join(submission.storage_dir, os.path.basename(source_path))
                if source_path != target_path:
                    if self.delete_files:
                        shutil.move(source_path, target_path)
                        logging.debug(f"moved file from {source_path} to {target_path}")
                    else:
                        shutil.copy2(source_path, target_path)
                        logging.debug(f"copied file from {source_path} to {target_path}")

                if isinstance(file_submission, tuple):
                    updated_files.append((target_path, dest_path))
                else:
                    updated_files.append(target_path)

            submission.files = updated_files

        except Exception as e:
            logging.error(f"I/O error moving files into {submission.storage_dir}: {e}")
            report_exception()

    @use_db(name='collection')
    def schedule_submission(self, submission, db, c):
        try:
            # add this as a workload item to the database queue
            work_id = execute_with_retry(db, c, self.insert_workload, (submission,), commit=True)
            logging.info(f"scheduled {submission.description} mode {submission.analysis_mode} work_id {work_id}")
            
        except Exception as e:
            # something went wrong -- delete our incoming directory if we created one
            target_dir = self.get_submission_target_dir(submission)
            if os.path.exists(target_dir):
                try:
                    shutil.rmtree(target_dir)
                except Exception as e:
                    logging.error("unable to delete directory {}: {}".format(target_dir, e))

            raise e

        self.submission_count += 1

    def cleanup_submission(self, submission):
        """Cleans up the given submission."""
        # if we are deleting files we add then we would have moved the file instead of copying it
        pass 

    def insert_workload(self, db, c, next_submission):
        c.execute("INSERT INTO incoming_workload ( type_id, mode, work ) VALUES ( %s, %s, %s )",
                 (self.workload_type_id, next_submission.analysis_mode, pickle.dumps(next_submission)))

        if c.lastrowid is None:
            raise RuntimeError("missing lastrowid for INSERT transaction")

        work_id = c.lastrowid

        # assign this work to each configured group
        node_groups = self.remote_node_groups
        
        # does this submission have a defined set to groups to send to?
        if next_submission.group_assignments:
            node_groups = [ng for ng in node_groups if ng.name in next_submission.group_assignments]

            if not node_groups:
                # default to all groups if we end up with an empty list
                logging.error(f"group assignment {next_submission.group_assignments} does not map to any known groups")
                node_groups = self.remote_node_groups
            
        for remote_node_group in node_groups:
            c.execute("INSERT INTO work_distribution ( work_id, group_id ) VALUES ( %s, %s )",
                     (work_id, remote_node_group.group_id))

        return work_id

    # subclasses can override this function to provide additional functionality
    def extended_collection(self):
        self.execute_in_loop(self.execute_extended_collection)

    def execute_in_loop(self, target):
        while True:
            if self.is_service_shutdown:
                return

            try:
                # refresh settings
                saq.settings.load()

                # delete expired persistent data every so often
                if time.time() - self.persistent_clear_time > self.service_config.getint("persistence_clear_seconds", 60):
                    expiration_timedelta = timedelta(seconds=self.service_config.getint("persistence_expiration_seconds", 24*60*60))
                    unmodified_expiration_timedelta = timedelta(seconds=self.service_config.getint("persistence_unmodified_expiration_seconds", 4*60*60))
                    self.delete_expired_persistent_keys(expiration_timedelta, unmodified_expiration_timedelta)
                    self.persistent_clear_time = time.time()

                wait_seconds = target()
                if wait_seconds is None:
                    wait_seconds = 1

                self.service_shutdown_event.wait(wait_seconds)
                continue
            except NotImplementedError:
                return
            except Exception as e:
                logging.error(f"unable to execute {target}: {e}")
                report_exception()
                self.service_shutdown_event.wait(1)
                continue
            finally:
                saq.db.remove()

    def execute_extended_collection(self):
        """Executes custom collection routines. 
           Returns the number of seconds to wait until the next time this should be called."""
        raise NotImplementedError()

    def debug_extended_collection(self):
        """Executes custom collection routines in debug mode. """
        return self.execute_extended_collection()

    def get_next_submission(self):
        """Returns the next Submission object to be submitted to the remote nodes."""
        try:
            return self.submission_list.get(block=True, timeout=1)
        except queue.Empty:
            return None

def parse_schedule(schedule):
    if schedule == "*":
        return { "interval": True, "value": 1 }
    elif schedule.startswith("*/"):
        return { "interval": True, "value": int(schedule[2:]) }
    else:
        return { "interval": False, "value": int(schedule) }
