#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import absolute_import

from sys import exc_info, version_info

from core import output
from core.config import CONFIG
from core.tools import geolocate

from geoip2.database import Reader

if version_info[0] < 3:
    import pymysql
    pymysql.install_as_MySQLdb()

try:
    from MySQLdb import (Error, OperationalError)
except ImportError:
    try:
        from MySQLdb._exceptions import (Error, OperationalError)
    except ImportError:
        from _mysql_exceptions import (Error, OperationalError) # type: ignore

from twisted.enterprise.adbapi import ConnectionPool
from twisted.python.compat import reraise
from twisted.python.log import msg


class ReconnectingConnectionPool(ConnectionPool):
    """
    Reconnecting adbapi connection pool for MySQL.
    This class improves on the solution posted at
    http://www.gelens.org/2008/09/12/reinitializing-twisted-connectionpool/
    by checking exceptions by error code and only disconnecting the current
    connection instead of all of them.
    Also see:
    http://twistedmatrix.com/pipermail/twisted-python/2009-July/020007.html
    """

    def _runInteraction(self, interaction, *args, **kw):

        def rerise_exception(conn):
            _, excValue, excTraceback = exc_info()
            try:
                conn.rollback()
            except Exception:
                msg('Rollback failed')
            reraise(excValue, excTraceback)

        conn = self.connectionFactory(self)
        trans = self.transactionFactory(self, conn)
        try:
            result = interaction(trans, *args, **kw)
            trans.close()
            conn.commit()
            return result
        except OperationalError as e:
            if e.args[0] not in (2003, 2006, 2013):
                rerise_exception(conn)
            else:
                conn = self.connections.get(self.threadID())
                self.disconnect(conn)
                return ConnectionPool._runInteraction(self, interaction, *args, **kw)
        except Exception:
            rerise_exception(conn)


class Output(output.Output):

    def local_log(self, message):
        if self.debug:
            msg(message)

    def start(self):
        host = CONFIG.get('output_mysql', 'host', fallback='localhost')
        database = CONFIG.get('output_mysql', 'database', fallback='siphoney')
        user = CONFIG.get('output_mysql', 'username', fallback='siphoney', raw=True)
        password = CONFIG.get('output_mysql', 'password', fallback='', raw=True)
        port = CONFIG.getint('output_mysql', 'port', fallback=3306)

        self.debug = CONFIG.getboolean('output_mysql', 'debug', fallback=False)
        self.geoip = CONFIG.getboolean('output_mysql', 'geoip', fallback=True)

        try:
            self.dbh = ReconnectingConnectionPool(
                'MySQLdb',
                host=host,
                db=database,
                user=user,
                passwd=password,
                port=port,
                charset='utf8',
                use_unicode=True,
                cp_min=1,
                cp_max=1
            )
        except Error as e:
            self.local_log('output_mysql: MySQL Error {}: "{}"'.format(e.args[0], e.args[1]))

        if self.geoip:
            geoipdb_city_path = CONFIG.get('output_mysql', 'geoip_citydb', fallback='data/GeoLite2-City.mmdb')
            geoipdb_asn_path = CONFIG.get('output_mysql', 'geoip_asndb', fallback='data/GeoLite2-ASN.mmdb')
            try:
                self.reader_city = Reader(geoipdb_city_path)
            except Exception:
                self.reader_city = None
                self.local_log('Failed to open City GeoIP database {}'.format(geoipdb_city_path))

            try:
                self.reader_asn = Reader(geoipdb_asn_path)
            except Exception:
                self.reader_asn = None
                self.local_log('Failed to open ASN GeoIP database {}'.format(geoipdb_asn_path))

    def stop(self):
        if self.geoip:
            if self.reader_city is not None:
                self.reader_city.close()
            if self.reader_asn is not None:
                self.reader_asn.close()

    def write(self, event):
        """
        TODO: Check if the type (date, datetime or timestamp) of columns is appropriate for your needs and timezone
        - MySQL Documentation - The DATE, DATETIME, and TIMESTAMP Types
            (https://dev.mysql.com/doc/refman/5.7/en/datetime.html):
        "MySQL converts TIMESTAMP values from the current time zone to UTC for storage,
        and back from UTC to the current time zone for retrieval.
        (This does not occur for other types such as DATETIME.)"
        """
        self.dbh.runInteraction(self.connect_event, event)

    def simple_query(self, txn, sql, args):
        if self.debug:
            if len(args):
                self.local_log("output_mysql: MySQL query: {} {}".format(sql, repr(args)))
            else:
                self.local_log("output_mysql: MySQL query: {}".format(sql))
        try:
            if len(args):
                txn.execute(sql, args)
            else:
                txn.execute(sql)
            result = txn.fetchall()
        except Exception as e:
            self.local_log('output_mysql: MySQL Error: {}'.format(e))
            result = None
        return result

    def get_id(self, txn, table, column, entry):
        r = self.simple_query(txn, "SELECT `id` FROM `{}` WHERE `{}` = %s".format(table, column), (entry, ))
        if r:
            id = r[0][0]
        else:
            self.simple_query(txn, "INSERT INTO `{}` (`{}`) VALUES (%s)".format(table, column), (entry, ))
            r = self.simple_query(txn, 'SELECT LAST_INSERT_ID()', ())
            if r:
                id = int(r[0][0])
            else:
                id = 0
        return id

    def get_request_id(self, txn, headers):
        accept_id = allow_id = callids_id = contact_id = conlen_id = contype_id = cseq_id = expires_id = from_id = forward_id = supported_id = to_id = agent_id = via_id = auth_id = 0
        for key, value in headers.items():
            if key == 'accept':
                accept_id = self.get_id(txn, 'accepts', 'value', value)
            elif key == 'allow':
                allow_id = self.get_id(txn, 'allows', 'value', value)
            elif key == 'call-id':
                callids_id = self.get_id(txn, 'callids', 'value', value)
            elif key == 'contact':
                contact_id = self.get_id(txn, 'contacts', 'value', value)
            elif key == 'content-length':
                conlen_id = self.get_id(txn, 'con_lens', 'value', value)
            elif key == 'content-type':
                contype_id = self.get_id(txn, 'con_types', 'value', value)
            elif key == 'cseq':
                cseq_id = self.get_id(txn, 'cseqs', 'value', value)
            elif key == 'expires':
                expires_id = self.get_id(txn, 'expires', 'value', value)
            elif key == 'from':
                from_id = self.get_id(txn, 'froms', 'value', value)
            elif key == 'max-forwards':
                forward_id = self.get_id(txn, 'forwards', 'value', value)
            elif key == 'supported':
                supported_id = self.get_id(txn, 'supporteds', 'value', value)
            elif key == 'to':
                to_id = self.get_id(txn, 'tos', 'value', value)
            elif key == 'user-agent':
                agent_id = self.get_id(txn, 'agents', 'value', value)
            elif key == 'via':
                via_id = self.get_id(txn, 'vias', 'value', value)
            elif key == 'authorization':
                auth_id = self.get_id(txn, 'authorizations', 'value', value)
            else:
                self.get_id(txn, 'unknown', 'header', key)
        self.simple_query(txn,
            """
            INSERT INTO `requests` (`accept`, `allow`, `callid`, `contact`, `con_len`, `con_type`, `cseq`, `expires`, `from`, `forwards`, `supported`, `to`, `agent`, `via`, `authorization`)
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            """,
            (accept_id, allow_id, callids_id, contact_id, conlen_id, contype_id, cseq_id, expires_id, from_id, forward_id, supported_id, to_id, agent_id, via_id, auth_id, )
        )
        r = self.simple_query(txn, 'SELECT LAST_INSERT_ID()', ())
        if r:
            id = int(r[0][0])
        else:
            id = 0
        return id

    def connect_event(self, txn, event):
        remote_ip = event['src_ip']
        sensor_id = self.get_id(txn, 'sensors', 'name', event['sensor'])
        message_id = self.get_id(txn, 'messages', 'value', event['message'])
        request_id = self.get_request_id(txn, event['headers'])

        self.simple_query(txn,
            """
            INSERT INTO `connections` (`timestamp`, `ip`, `remote_port`, `local_host`, `local_port`, `request`, `message`, `sensor`)
            VALUES (FROM_UNIXTIME(%s), %s, %s, %s, %s, %s, %s, %s)
            """,
            (event['unixtime'], remote_ip, event['src_port'], event['dst_ip'],
             event['dst_port'], request_id, message_id, sensor_id, ))

        if self.geoip:
            country, country_code, city, org, asn_num = geolocate(remote_ip, self.reader_city, self.reader_asn)
            self.simple_query(txn, """
                INSERT INTO `geolocation` (`ip`, `country_name`, `country_iso_code`, `city_name`, `org`, `org_asn`)
                VALUES (%s, %s, %s, %s, %s, %s)
                ON DUPLICATE KEY UPDATE
                    `country_name` = %s,
                    `country_iso_code` = %s,
                    `city_name` = %s,
                    `org` = %s,
                    `org_asn` = %s
                """,
                (remote_ip, country, country_code, city, org, asn_num, country, country_code, city, org, asn_num, ))
