
from kazoo.client import KazooClient
from kafka.coordinator.assignors.roundrobin import RoundRobinPartitionAssignor
from kafka.coordinator.assignors.range import RangePartitionAssignor
from kafka.common import UnknownTopicOrPartitionError, KafkaTimeoutError, KafkaUnavailableError, LeaderNotAvailableError, NotLeaderForPartitionError, ReplicaNotAvailableError, ConnectionError, KafkaConfigurationError, FailedPayloadsError, OffsetOutOfRangeError
from kafka import KafkaConsumer
from kafka.client import KafkaClient
import os
import sys
import string
import time
import datetime
import signal
import errno
import base64
import math
import socket
import subprocess
import threading
import tempfile
import argparse
from multiprocessing import Process, Value
import gzip
import psutil

import json
import fcntl
import re

from ssl import SSLError

from dateutil.relativedelta import relativedelta
from pwd import getpwnam
from grp import getgrnam

import traceback

from socket import error as SocketError
import requests

import urllib

from requests.packages import urllib3
from requests.models import Response

from hawkSoar.common.logger import get_logger
from hawkSoar.common.process import ExceptionSafeProcess
from hawkSoar.plugins.base import BasePlugin
from hawkSoar.common.errors import HawkIOException, HawkBaseError


urllib3.disable_warnings()


translationTable = {
    "_id": "event_id",
    "w": "weight",
    "a_name": "alert_name",
    "pri": "priority",
    "at_name": "alert_category",
    "c": "count",
    "b": "blocked",
    "ip_s": "ip_src",
    "ip_d": "ip_dst",
    "ip_sh": "ip_src_host",
    "ip_dh": "ip_dst_host",
    "correlation_username": "username_src",
    "target_username": "username_dst",
    "g_n": "group_name",
    "c_key": "class_type",
    "c_name": "class_name",
    "o_name": "os_name",
    "s_g_cc2": "ip_src_geoip_cc2",
    "s_g_n": "ip_src_geoip_name",
    "s_g_c": "ip_src_geoip_city",
    "s_g_r": "ip_src_geoip_region",
    "s_g_la": "ip_src_geoip_latitude",
    "s_g_lo": "ip_src_geoip_longitude",
    "d_g_cc2": "ip_dst_geoip_cc2",
    "d_g_n": "ip_dst_geoip_name",
    "d_g_c": "ip_dst_geoip_city",
    "d_g_r": "ip_dst_geoip_region",
    "d_g_la": "ip_dst_geoip_latitude",
    "d_g_lo": "ip_dst_geoip_longitude",
    "r_name": "resource_name",
    "r_addr": "resource_address",
    "rs_n": "result_name",
    "rs_v": "result_value"
}


class IngestPlugin(BasePlugin):

    ID = 'hawk-v-stream'
    NAME = 'HAWK-VStream'
    DESCRIPTION = 'Consume data from HAWK VStream data store'
    TYPE = 'ingest'
    DEPENDENCIES = ['python-kafka']
    VERSION = '1.0.0'
    MODULE = 'playbook.ingest.Events'
    ENABLED = True
    Exceptions = []
    process_list = []
    active_sessions = {}

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

    def process_task(self, task):
        if task['cmd'] == 'addPlaybookStep':
            self._add_playbook_step(task['data'])

            # start up new worker when new task arrives.
            # print(task)
            # print(task['data'])
            if not task['data']['id'] in self.active_sessions:
                self.active_sessions[task['data']['id']] = []

            if not task['data']['step']['id'] in self.active_sessions[task['data']['id']]:
                # start up worker

                vdb = VStream()

                if 'tags' in task['data']['step']:
                    tags = task['data']['step']['tags']
                else:
                    tags = []

                if 'tags' in task['data']['playbook'] and isinstance(task['data']['playbook']['tags'], list):
                    for t in task['data']['playbook']['tags']:
                        if not t in tags:
                            tags.append(t)

                vdb.setConfig(self.ID, task['data']['playbook']['group'], self.VERSION, self.MODULE, task['data']['id'], task['data']['step']['id'],
                              task['data']['step']['next'], tags, task['data']['config'], self.outgoing_queue)
                vdb.start()
                self.process_list.append(vdb)

                self.active_sessions[task['data']['id']].append(
                    task['data']['step']['id'])

        else:
            # do the work
            self.logger.info("Unknown task: %s" % json.dumps(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):
        self.logger.info('HAWK V-Stream plugin is now shutting down.')
        self.logger.info(
            'HAWK V-Stream plugin is waiting for process list to terminate.')
        for process in self.process_list:
            process.shutdown()
            process.join()

        # super().shutdown()
        self.logger.info('HAWK V-Stream shutdown complete.')


class VStream(Process):

    def setConfig(self, pluginId, groupId, version, module, playbookId, taskId, nextStep, tags, kafkaConfig, outgoing_queue):
        self.groupId = groupId
        self.pluginId = pluginId
        self.version = version
        self.module = module
        self.logger = get_logger()
        self.stop_condition = Value('i', 0)
        # hawk only uses zookeeper
        self.playbookId = playbookId
        self.taskId = taskId
        self.nextStep = nextStep
        self.config = kafkaConfig
        self.tags = tags
        self.outgoing_queue = outgoing_queue

    def shutdown(self):
        self.stop_condition.value = 1

    def run(self):

        self.logger.info("VStream: starting consumer for: %s" %
                         self.config['hosts'])
        self.logger.info(json.dumps(self.config))

        self.kafka_nodes = self._get_bootstrap_servers(self.config['hosts'])

        self.logger.info(
            "VStream: dynamically gathered data nodes: %s" % self.kafka_nodes)

        if 'timeout' in self.config:
            timeout = float(self.config['timeout'])
            if timeout < 1000:
                timeout = timeout * 1000  # convert to ms
        else:
            # 30 sec
            timeout = 30000

        self.consumer = KafkaConsumer(bootstrap_servers=self.kafka_nodes, client_id='hawk-soard', group_id=str(
            self.config['consumer_group']), value_deserializer=lambda m: json.loads(m.decode('utf-8')), consumer_timeout_ms=timeout)
        # self.consumer = KafkaConsumer(bootstrap_servers=self.kafka_nodes, client_id='hawk-soard', group_id="%s_%s" % (str(
        #    self.playbookId), str(self.taskId)), value_deserializer=lambda m: json.loads(m.decode('utf-8')), consumer_timeout_ms=timeout)

        topic_regex = re.compile(self.config['topic_regex'])
        if not topic_regex:
            topic_regex = re.compile('^events\.')

        self.logger.info("VStream: listening on topic regex: %s" %
                         str(topic_regex))

        x = self.consumer.subscribe((), topic_regex)

        # self.logger.debug("V-Stream Topics: %s" % ( self.consumer.topics().join(', ') ) )

        while self.stop_condition.value != 1:
            for message in self.consumer:
                msg = message.value
                # jobs done!
                for key in translationTable.keys():
                    # print("Checking if %s exists" % key)
                    if key in msg:
                        # print("%s exists" % key)
                        # print("mapping to %s " % translationTable[key] )
                        msg[translationTable[key]] = msg[key]
                        del msg[key]

                out_msg = {"cmd": "task", "route": "publish",
                           "data": {
                               "playbookId": self.playbookId, "step": self.nextStep, "group": self.groupId,
                               "data": {
                                   "event": msg,
                                   "artifacts": [
                                       {"plugin_id": self.pluginId, "name": "version",
                                        "module": self.module, "type": "app.version",  "value": self.version}
                                   ],
                                   "tags": self.tags,
                                   "return_value": True
                               }
                           }
                           }
                self.outgoing_queue.put(out_msg)

    def _get_bootstrap_servers(self, zservers):

        # Dynamic broker configuration.
        try:
            ZooKeeper = KazooClient(hosts=zservers)
            ZooKeeper.start()
        except Exception as e:
            errstr = 'Zookeeper Failed to connect.'
            errors = getattr(e, 'errors', None)
            # self.logger.exception(errstr)
            raise e
            # raise HawkZookeeperError(errstr, errors)

        try:
            broker_ids = ZooKeeper.get_children('/brokers/ids')
        except Exception as e:
            errstr = 'Zookeeper Failed to get childern'
            errors = getattr(e, 'errors', None)
            # self.logger.exception(errstr)
            raise e
            # raise HawkZookeeperError(errstr, errors)

        kafka_servers = []

        for bid in broker_ids:
            try:
                results_str, stats = ZooKeeper.get('/brokers/ids/%s' % (bid))
                results = json.loads(results_str)
                kafka_servers.append('%s:%r' %
                                     (results['host'], int(results['port'])))
            except Exception as e:
                errstr = 'Zookeeper failed to get childern details'
                errors = getattr(e, 'errors', None)
                # self.logger.exception(errstr)
                # raise HawkZookeeperError(errstr, errors)
                raise e

        ZooKeeper.stop()

        return ','.join(kafka_servers)


# class HawkZookeeperError(HawkIOException):
#    """Class to catch hawk core errors. """
#
#    def __init__(self, msg, errors=None):
#        HawkBaseError.__init__(self, msg, errors)
