import datetime
import dateutil.parser
import dateutil.tz
import logging
import pytz
import re
import subprocess

from cbapi.response import *
from cbapi.psc import threathunter

from cbinterface.cli import load_configured_environments

#from cbinterface.modules.process import ProcessWrapper
#from cbinterface.modules.helpers import as_configured_timezone
from cbinterface.config import get_timezone

from lib.modules.DetectionModule import *

class Module(DetectionModule):
    def __init__(self, name, event_json):

        super().__init__(name=name, event_json=event_json)

    def run(self):
        self.logger.debug('Running the {} detection module'.format(self.name))

        # These are the companies that will get cbinterface queries.
        ignore_these_companies = self.config['ignore_these_companies']
        company_names = set()
        for alert in self.event_json['ace_alerts']:
            if alert['company_name'] and not alert['company_name'] in ignore_these_companies:
                company_names.add(alert['company_name'])

        # These are the process names that will be ignored from the queries.
        ignore_these_process_names = self.config['ignore_these_process_names']
        ignore_these_process_names_string = ''
        for process_name in ignore_these_process_names:
            ignore_these_process_names_string += '-process_name:{} '.format(process_name)

        # These are common file names that should be ignored in queries.
        ignore_these_filenames = self.config['ignore_these_filenames']

        # Ignore these MD5s all together
        ignore_these_md5s = self.config['ignore_these_md5s']

        # These are the process MD5s that will be ignored from the queries.
        ignore_these_process_md5s = self.config['ignore_these_process_md5s']
        ignore_these_process_md5s_string = ''
        for process_md5 in ignore_these_process_md5s:
            ignore_these_process_md5s_string += '-process_md5:{} '.format(process_md5)

        # These are the computer names that will be ignored from the queries.
        ignore_these_computers = self.config['ignore_these_computers']
        ignore_these_computers_string = ''
        for computer in ignore_these_computers:
            ignore_these_computers_string += '-hostname:{} '.format(computer)

        # These are the users that will be ignored from the queries.
        ignore_these_users = self.config['ignore_these_users']
        ignore_these_users_string = ''
        for user in ignore_these_users:
            ignore_these_users_string += '-username:{} '.format(user)


        # Get all of the good Windows - FileName indicators from the event.
        good_indicators = [i for i in self.event_json['indicators'] if not i['whitelisted']]
        filenames = list(set([i['value'] for i in good_indicators if i['type'] == 'Windows - FileName' and (i['status'] == 'New' or i['status'] == 'Analyzed' or i['status'] == 'In Progress')]))
        # remove filenames we want to ignore
        filenames = [f for f in filenames if f not in ignore_these_filenames]
        self.logger.debug('cbinterface filenames: {}'.format(filenames))

        # Get all of the good Hash - MD5 indicators from the event.
        md5s = list(set([i['value'] for i in good_indicators if i['type'] == 'Hash - MD5' and (i['status'] == 'New' or i['status'] == 'Analyzed' or i['status'] == 'In Progress')]))
        # remove and md5s we want to ignore
        md5s = [md5 for md5 in md5s if md5 not in ignore_these_md5s]
        self.logger.debug('cbinterface md5s: {}'.format(md5s))

        # Get all of the good Hash - SHA256 indicators from the event.
        # only use CBC/PSC with
        sha256s = list(set([i['value'] for i in good_indicators if i['type'] == 'Hash - SHA256' and (i['status'] == 'New' or i['status'] == 'Analyzed' or i['status'] == 'In Progress')]))
        self.logger.debug('cbinterface md5s: {}'.format(sha256s))

        # Get the start time.
        start_time = ''
        if self.event_json['emails']:
            start_time = self.event_json['emails'][0]['received_time']
            self.logger.debug("Using start time from emails received_time: {}".format(start_time))
        elif self.event_json['ace_alerts']:
            start_time = self.event_json['ace_alerts'][0]['time']
            self.logger.debug("Using start time from ace alert time: {}".format(start_time))

        # We need to make sure the start time is in the format "YYYY-MM-DD", which is 10 characters long.
        start_time = start_time[0:10]

        # Force the start time to begin at 00:00:00.
        start_time = '{} 00:00:00'.format(start_time)
        start_time = datetime.datetime.strptime(start_time, '%Y-%m-%d %H:%M:%S')

        # Continue if we have a start time.
        if start_time:

            cb = None
            environment_map = load_configured_environments()
            for product,environments in environment_map.items():
                for profile in environments:
                    if product == "response":
                        from cbinterface.response.query import make_process_query
                        from cbinterface.response.process import print_process_info
                        cb = CbResponseAPI(profile=profile)
                    elif product == "psc":
                        from cbinterface.psc.query import make_process_query, convert_from_legacy_query
                        from cbinterface.psc.process import print_process_info
                        cb = threathunter.CbThreatHunterAPI(profile=profile)

                    for filename in filenames:
                        cli_friendly_filename = re.sub(r"(!|\$|#|&|\"|\'|\(|\)|\||<|>|`|\\\|;)", r"\\\1", filename)
                        query = '{} {} {} {} (filemod:"{}" OR cmdline:"{}")'.format(ignore_these_process_names_string, ignore_these_process_md5s_string, ignore_these_computers_string, ignore_these_users_string, cli_friendly_filename, cli_friendly_filename)
                        if product == "psc":
                            #query = '{} {} {} {} (filemod:"{}" OR cmdline:"{}")'.format(ignore_these_process_names_string, ignore_these_process_md5s_string, ignore_these_computers_string, ignore_these_users_string, cli_friendly_filename, cli_friendly_filename)
                            query = convert_from_legacy_query(cb, query)
   
                        self.logger.info("Searching Carbon Black data for file modifications with: {}".format(query))
                        processes = []
                        try:
                            processes = make_process_query(cb, query, start_time=start_time)
                        except Exception as e:
                            self.logger.exception(f"problem querying {product}:{profile}")
                        if not processes:
                            continue

                        if len(processes) > 100:
                            self.logger.warning("{} processes returned by Carbon Black for '{}' since '{}'".format(len(processes), query, start_time.isoformat()))
                            self.detections.append('! DETECTED {} PROCESSSES THAT MADE "{}" LIKE FILEMODS ! -- TOO LARGE of dataset to display, only including first 100 results in extra detections.'.format(len(processes), filename))
                            self.tags.append('incidents')
                            self.tags.append('exploitation')
                            for proc in processes[:100]:
                                hostname = proc.get('hostname') or proc.get('device_name') or "unknown"
                                username = proc.get('username') or proc.get('process_username') or "unknown"
                                if isinstance(username, list):
                                    username = username[0]
                                process_name = proc.get('process_name')
                                cmdline = proc.get('cmdline') or proc.get('process_cmdline') or "None"
                                if proc.get('segment_id'):
                                    # response
                                    webui_link = f"{cb.credentials.url}/#analyze/{proc.get('id')}/{proc.get('segment_id')}"
                                else:
                                    webui_link = f"{cb.credentials.url}/analyze?processGUID={proc.get('process_guid')}"
                                proc_summary = "Hostname:{} User:{} Process_Name:{} Command_Line:{} GUI_Link:{}".format(hostname,
                                                                                                                        username,
                                                                                                                        process_name,
                                                                                                                        cmdline,
                                                                                                                        webui_link)
                                self.extra.append(proc_summary)
                        else:
                            for proc in processes:
                                if proc.get('segment_id'):
                                    # response
                                    webui_link = f"{cb.credentials.url}/#analyze/{proc.get('id')}/{proc.get('segment_id')}"
                                else:
                                    webui_link = f"{cb.credentials.url}/analyze?processGUID={proc.get('process_guid')}"
                                self.detections.append('! DETECTED FILENAME {} ! {}'.format(filename, webui_link))
                                self.tags.append('incidents')
                                self.tags.append('exploitation')
                                self.extra.append(print_process_info(proc, return_string=True))

                    for md5 in md5s:
                        query = '{} {} {} {} md5:{}'.format(ignore_these_process_names_string, ignore_these_process_md5s_string, ignore_these_computers_string, ignore_these_users_string, md5)
                        if product == "psc":
                            #query = '{} {} {} {} hash:{}'.format(ignore_these_process_names_string, ignore_these_process_md5s_string, ignore_these_computers_string, ignore_these_users_string, md5)
                            query = convert_from_legacy_query(cb, query)
   
                        self.logger.info("Searching Carbon Black data for MD5s with: {}".format(query))
                        processes = []
                        try:
                            processes = make_process_query(cb, query, start_time=start_time)
                        except Exception as e:
                            self.logger.exception(f"problem querying {product}:{profile}")
                        if not processes:
                            continue

                        if len(processes) > 100:
                            self.logger.warning("{} processes returned by Carbon Black for '{}' since '{}'".format(len(processes), query, start_time.isoformat()))
                            self.detections.append('! DETECTED {} PROCESSSES THAT MADE WITH MD5 {} ! -- TOO LARGE of dataset to display, only including first 100 results in extra detections.'.format(len(processes), md5))
                            self.tags.append('incidents')
                            self.tags.append('exploitation')
                            for proc in processes[:100]:
                                hostname = proc.get('hostname') or proc.get('device_name') or "unknown"
                                username = proc.get('username') or proc.get('process_username') or "unknown"
                                if isinstance(username, list):
                                    username = username[0]
                                process_name = proc.get('process_name')
                                cmdline = proc.get('cmdline') or proc.get('process_cmdline') or "None"
                                if proc.get('segment_id'):
                                    # response
                                    webui_link = f"{cb.credentials.url}/#analyze/{proc.get('id')}/{proc.get('segment_id')}"
                                else:
                                    webui_link = f"{cb.credentials.url}/analyze?processGUID={proc.get('process_guid')}"

                                proc_summary = "Hostname:{} User:{} Process_Name:{} Command_Line:{} GUI_Link:{}".format(hostname,
                                                                                                                        username,
                                                                                                                        process_name,
                                                                                                                        cmdline,
                                                                                                                        webui_link)
                                self.extra.append(proc_summary)
                        else:
                            for proc in processes:
                                if proc.get('segment_id'):
                                    # response
                                    webui_link = f"{cb.credentials.url}/#analyze/{proc.get('id')}/{proc.get('segment_id')}"
                                else:
                                    webui_link = f"{cb.credentials.url}/analyze?processGUID={proc.get('process_guid')}"
                                self.detections.append('! DETECTED MD5 {} ! {}'.format(md5, webui_link))
                                self.tags.append('incidents')
                                self.tags.append('exploitation')
                                # just using the convienient display of ProcessWrapper
                                self.extra.append(print_process_info(proc, return_string=True))

                    for sha256 in sha256s:
                        if product == "response":
                            break
                        query = f'{ignore_these_process_names_string} {ignore_these_process_md5s_string} {ignore_these_computers_string} {ignore_these_users_string}'
                        query = convert_from_legacy_query(cb, query)
                        query += f" hash:{sha256}"
   
                        self.logger.info("Searching Carbon Black data for SHA256 with: {}".format(query))
                        processes = []
                        try:
                            processes = make_process_query(cb, query, start_time=start_time)
                        except Exception as e:
                            self.logger.exception(f"problem querying {product}:{profile}")
                        if not processes:
                            continue

                        if len(processes) > 100:
                            self.logger.warning("{} processes returned by Carbon Black for '{}' since '{}'".format(len(processes), query, start_time.isoformat()))
                            self.detections.append('! DETECTED {} PROCESSS HITS on SHA256 {} ! -- TOO LARGE of dataset to display, only including first 100 results in extra detections.'.format(len(processes), sha256))
                            self.tags.append('incidents')
                            self.tags.append('exploitation')
                            for proc in processes[:100]:
                                hostname = proc.get('hostname') or proc.get('device_name') or "unknown"
                                username = proc.get('username') or proc.get('process_username') or "unknown"
                                if isinstance(username, list):
                                    username = username[0]
                                process_name = proc.get('process_name')
                                cmdline = proc.get('cmdline') or proc.get('process_cmdline') or "None"
                                if proc.get('segment_id'):
                                    # response
                                    webui_link = f"{cb.credentials.url}/#analyze/{proc.get('id')}/{proc.get('segment_id')}"
                                else:
                                    webui_link = f"{cb.credentials.url}/analyze?processGUID={proc.get('process_guid')}"

                                proc_summary = "Hostname:{} User:{} Process_Name:{} Command_Line:{} GUI_Link:{}".format(hostname,
                                                                                                                        username,
                                                                                                                        process_name,
                                                                                                                        cmdline,
                                                                                                                        webui_link)
                                self.extra.append(proc_summary)
                        else:
                            for proc in processes:
                                if proc.get('segment_id'):
                                    # response
                                    webui_link = f"{cb.credentials.url}/#analyze/{proc.get('id')}/{proc.get('segment_id')}"
                                else:
                                    webui_link = f"{cb.credentials.url}/analyze?processGUID={proc.get('process_guid')}"
                                self.detections.append('! DETECTED SHA256 {} ! {}'.format(sha256, webui_link))
                                self.tags.append('incidents')
                                self.tags.append('exploitation')
                                self.extra.append(print_process_info(proc, return_string=True))
