

import os
import sys
import json
import requests
import time
from urllib.parse import urlparse

import xml
import radix
import csv
import zipfile
import pickle
import pysos

import shutil

from hawkSoar.common.logger import get_logger
from hawkSoar.common.process import ExceptionSafeProcess
from hawkSoar.plugins.base import BasePlugin
from hawkSoar.lib.query_parser import QueryInterpreter


class InvestigatePlugin(BasePlugin):

    ID = 'hawk-data-enrichment'
    NAME = 'Data Enrichment'
    DESCRIPTION = 'Enrich the artifacts with external data available as XML, CSV, or JSON'
    TYPE = 'investigate'
    DEPENDENCIES = ['py-radix', 'xml', 'pysos']
    VERSION = '0.0.1'
    MODULE = "playbook.investigate.DataEnrichment"
    ENABLED = True
    Exceptions = []
    playbookCache = {}
    activeRequests = []
    activeURL = {}

    def __init__(self, plugin_administrator, config=None):
        super().__init__(plugin_administrator, config=config, plugin_path=__file__)
        self.config = config
        self.logger = get_logger()
        self.playbookCache = {}

    def download_file(self, url, filepath, timeout=30, retries=5, verify=True):

        with requests.Session() as c:
            while retries > 0:
                try:
                    self.logger.debug(
                        "Downloading file path: %s url: %s" % (filepath, url))
                    res = c.get(url, verify=True, timeout=timeout)
                    with open(filepath, 'wb') as f:
                        f.write(res.content)
                    return True
                except Exception as e:
                    retries -= 1

            return False

    def process_file(self, filename):

        if zipfile.is_zipfile(filename):
            with zipfile.ZipFile(filename) as myzip:
                names = myzip.namelist()
                print(names)
                with myzip.open(names[0]) as f:
                    with open(filename[:-4], 'wb') as out:
                        out.write(f.read())

                        return filename[:-4]

        return filename

    def process_task(self, task):
        self.logger.debug("DATA ENRICHMENT TASK RECEIVED: %s" % task['cmd'])
        self.logger.debug(json.dumps(task))
        if task['cmd'] == 'addPlaybookStep':
            self._add_playbook_step(task['data'])
            # start up new worker when new task arrives.

            if not task['data']['id'] in self.playbookCache:
                self.playbookCache[task['data']['id']] = {}

            self.logger.debug("MATCH NEW STEP TASK")
            playbookId = task['data']['id']
            stepId = task['data']['step']['id']
            self.logger.debug(task['data']['step'])
            self.playbookCache[playbookId][stepId] = {
                'config':  task['data']['config'],
                'step':  task['data']['step']
            }

            # load indexed list from disk for lookups

            url = urlparse(
                self.playbookCache[playbookId][stepId]['config']['url'])
            if not url:
                raise Exception("Unable to parse url: %s" %
                                self.playbookCache[playbookId][stepId]['config']['url'])

            filename = self.config['plugin_data_path'] + \
                ('/%s/%s/' % (self.TYPE, self.ID)) + os.path.basename(url.path)
            filename = self.process_file(filename)

            if self.playbookCache[playbookId][stepId]['config']['index_type'] == 'text':
                dataType = 'text'
                """
                index_file = "%s.tdx" % filename
                if os.path.isfile(index_file):
                    textLookup = pysos.load(index_file)
                else:
                    textLookup = pysos.Dict(index_file)
                """
            elif self.playbookCache[playbookId][stepId]['config']['index_type'] == 'cidr':
                dataType = 'cidr'
                index_file = "%s.ndx" % filename
                cidrLookup = radix.Radix()
                # example rnode = rtree.add("10.0.0.0/8")
                # rnode.data["blah"] = "whatever you want"
            else:
                raise Exception("Unknown index type provided by configuration: %s" %
                                self.playbookCache[playbookId][stepId]['config']['index_type'])

            if not os.path.isfile(filename):
                # fetch 1st time
                # and process
                # self.
                if not self.download_file(self.playbookCache[playbookId][stepId]['config']['url'], filename, int(self.playbookCache[playbookId][stepId]['config']['timeout']), int(self.playbookCache[playbookId][stepId]['config']['retry'])):
                    raise Exception("Unable to download file, ignoring step request: %s" %
                                    self.playbookCache[playbookId][stepId]['config']['url'])

            else:
                # check date of file updated, and if expired, fetch and process
                mt = os.path.getmtime(filename)
                if (mt + (self.playbookCache[playbookId][stepId]['config']['refresh_hours'] * (3600)) < time.time()):
                    # replacing
                    self.logger.info("Updating (outdated) file: %s" % filename)
                    # remove index
                    try:
                        if os.path.isdir(index_file):
                            shutil.rmtree(self.path, ignore_errors=True)
                        else:
                            os.unlink(index_file)
                    except:
                        pass

                    if not self.download_file(self.playbookCache[playbookId][stepId]['config']['url'], filename, int(self.playbookCache[playbookId][stepId]['config']['timeout']), int(self.playbookCache[playbookId][stepId]['config']['retry'])):
                        raise Exception("Unable to download file, ignoring step request: %s" %
                                        self.playbookCache[playbookId][stepId]['config']['url'])

            if dataType == 'text':
                # load text index
                index_file = "%s.tdx" % filename
                if not os.path.isfile(index_file):
                    # now process the file
                    self.logger.info(
                        "Building index (%s) from local file: %s" % (dataType, filename))
                    with open(filename, 'r', encoding="utf-8-sig") as f:
                        if self.playbookCache[playbookId][stepId]['config']['format_type'] == 'csv':
                            # build csv index
                            if len(self.playbookCache[playbookId][stepId]['config']['target_headers']) > 0:
                                headers = self.playbookCache[playbookId][stepId]['config']['target_headers']
                                reader = csv.DictReader(f, fieldnames=headers)
                            else:
                                headers = []
                                header_reader = csv.reader(f)
                                for h in header_reader:
                                    headers = h
                                    break
                                reader = csv.DictReader(f, fieldnames=headers)

                            entries = 0
                            print("Loading entries into index")
                            textLookup = pysos.Dict(index_file)
                            for row in reader:
                                # print(row)
                                textLookup[row[self.playbookCache[playbookId]
                                               [stepId]['config']['target_lookup']]] = row
                                entries += 1
                            print("Loaded %d entries into index" % entries)

                            # textLookup.close()
                            #textLookup = pysos.load(index_file)
                            self.logger.debug(
                                "Index total number of %d records" % len(textLookup))
                            self.playbookCache[playbookId][stepId]['index'] = textLookup
                            self.playbookCache[playbookId][stepId]['type'] = dataType

                        elif self.playbookCache[playbookId][stepId]['config']['format_type'] == 'xml':
                            raise NotImplementedError
                        elif self.playbookCache[playbookId][stepId]['config']['format_type'] == 'json':
                            raise NotImplementedError
                else:
                    textLookup = pysos.load(index_file)
                    self.playbookCache[playbookId][stepId]['index'] = textLookup
                    self.playbookCache[playbookId][stepId]['type'] = dataType

            elif dataType == 'cidr':
                # cidrLookup
                index_file = "%s.ndx" % filename
                if not os.path.isfile(index_file):
                    self.logger.info(
                        "Building index (%s) from local file: %s" % (dataType, filename))
                    with open(filename, 'r', encoding="utf-8-sig") as f:
                        if self.playbookCache[playbookId][stepId]['config']['format_type'] == 'csv':

                            # build csv index
                            if len(self.playbookCache[playbookId][stepId]['config']['target_headers']) > 0:
                                headers = self.playbookCache[playbookId][stepId]['config']['target_headers']
                                reader = csv.DictReader(f, fieldnames=headers)
                            else:
                                headers = []
                                header_reader = csv.reader(f)
                                for h in header_reader:
                                    headers = h
                                    break
                                reader = csv.DictReader(f, fieldnames=headers)

                            for row in reader:
                                rnode = cidrLookup.add(
                                    row[self.playbookCache[playbookId][stepId]['config']['target_lookup']])
                                rnode.data = row

                            self.playbookCache[playbookId][stepId]['index'] = cidrLookup
                            self.playbookCache[playbookId][stepId]['type'] = dataType

                            with open(index_file, 'wb') as f:
                                pickle.dump(cidrLookup, f,
                                            pickle.HIGHEST_PROTOCOL)

                        elif self.playbookCache[playbookId][stepId]['config']['format_type'] == 'xml':
                            raise NotImplementedError
                        elif self.playbookCache[playbookId][stepId]['config']['format_type'] == 'json':
                            raise NotImplementedError

                else:
                    # pickle load
                    self.logger.info(
                        "Loading index (%s) from local file: %s" % (dataType, filename))
                    cidrLookup = pickle.load(open(index_file, "rb"))
                    self.playbookCache[playbookId][stepId]['index'] = cidrLookup
                    self.playbookCache[playbookId][stepId]['type'] = dataType

            self.logger.info(
                "Loading index (%s) completed." % (index_file))

        else:

            #self.logger.debug("DATA ENRICHMENT DO IT TASK")
            # self.logger.debug(task)
            playbookId = task['data']['playbookId']
            stepId = task['data']['step']

            step = self.playbookCache[playbookId][stepId]

            # print(self.get_step_by_id(task['data']['playbookId'], task['data']['step']))

            task['data']['data']['artifacts'].append(
                {"plugin_id": self.ID, "name": "version", "module": self.MODULE, "type": "app.version",  "value": self.VERSION})

            result = None

            stepObj = self.get_step_by_id(
                task['data']['playbookId'], task['data']['step'])

            for column in step['config']['lookup']:
                if column in task['data']['data'][step['config']['class']]:
                    val = task['data']['data'][step['config']['class']][column]
                    print("LOOKUP ENRICHMENT RECORD FOR value: %s" % str(val))
                    if self.playbookCache[playbookId][stepId]['config']['index_type'] == 'text':
                        try:
                            result = step['index'][val.strip()]
                        except KeyError as e:
                            result = None
                    elif self.playbookCache[playbookId][stepId]['config']['index_type'] == 'cidr':
                        res = step['index'].search_exact(val)
                        print(res)
                        if res and hasattr(res, 'data'):
                            result = res.data
                    else:
                        raise Exception(
                            "Unknown index type: %s" % self.playbookCache[playbookId][stepId]['config']['index_type'])

                if result:
                    # add data
                    print("FOUND ENRICHMENT DATA")
                    # print(result)
                    task['data']['data']['artifacts'].append(
                        {"plugin_id": self.ID, "type": "data.enrichment", "module": self.MODULE, "name": "%s.%s" % (stepObj['config']['class'], column), "value": result})

                    if 'tags' in stepObj:
                        for tag in stepObj['tags']:
                            if tag not in task['data']['data']['tags']:
                                task['data']['data']['tags'].append(tag)

                result = None

            task['data']['data']['return_value'] = True

            task['route'] = 'publish'

            # print(step)
            task['data']['step'] = step['step']['next']

            #self.logger.debug("OUTGOING ENRICHMENT TASK")
            # self.logger.debug(json.dumps(task))

            self.outgoing_queue.put(task)

    def add_exception(self, msg):
        self.Exceptions.append(msg)

    def check_exceptions(self):
        # if exceptiosn found, then log out
        return self.Exceptions

    def shutdown(self):
        # stop all on-going workers here
        self.logger.debug('Data enrichment plugin is now shutting down.')

        # stop all on-going workers here
        # super().shutdown()

        self.logger.debug('Data enrichment shutdown complete.')
