
from __future__ import absolute_import

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

from geoip2.database import Reader
from sqlite3 import OperationalError, IntegrityError

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


class Output(output.Output):

    def start(self):
        db_name = CONFIG.get('output_sqlite', 'db_file', fallback='data/siphoney.db')
        self.debug = CONFIG.getboolean('output_sqlite', 'debug', fallback=False)
        self.geoip = CONFIG.getboolean('output_sqlite', 'geoip', fallback=True)

        try:
            self.db = ConnectionPool(
                'sqlite3',
                database=db_name,
                check_same_thread=False
            )
        except OperationalError as e:
            msg(e)

        self.db.start()

        if self.geoip:
            geoipdb_city_path = CONFIG.get('output_sqlite', 'geoip_citydb', fallback='data/GeoLite2-City.mmdb')
            geoipdb_asn_path = CONFIG.get('output_sqlite', 'geoip_asndb', fallback='data/GeoLite2-ASN.mmdb')

            try:
                self.reader_city = Reader(geoipdb_city_path)
            except Exception:
                self.reader_city = None
                msg('Failed to open City GeoIP database {}'.format(geoipdb_city_path))
            try:
                self.reader_asn = Reader(geoipdb_asn_path)
            except Exception:
                self.reader_asn = None
                msg('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):
        self.db.runInteraction(self.connect_event, event)

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

    def get_id(self, txn, table, column, entry):
        r = self.simple_query(txn, "SELECT `id` FROM `{}` WHERE `{}` = ?".format(table, column), (entry, ))
        if r:
            id = r[0][0]
        else:
            self.simple_query(txn, "INSERT INTO `{}` (`{}`) VALUES (?)".format(table, column), (entry, ))
            r = self.simple_query(txn, 'SELECT LAST_INSERT_ROWID()', ())
            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 (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (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_ROWID()', ())
        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 (DATETIME(?, 'unixepoch', 'localtime'), ?, ?, ?, ?, ?, ?, ?)
            """,
            (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)
            result = self.simple_query(txn,
                """
                INSERT INTO `geolocation` (`ip`, `country_name`, `country_iso_code`, `city_name`, `org`, `org_asn`)
                VALUES (?, ?, ?, ?, ?, ?)
                """,
                (remote_ip, country, country_code, city, org, asn_num, ))
            if result is None:
                self.simple_query(txn,
                    """
                    UPDATE `geolocation` SET
                        `country_name` = ?,
                        `country_iso_code` = ?,
                        `city_name` = ?,
                        `org` = ?,
                        `org_asn` = ?
                    WHERE `ip` = ?
                    """,
                    (country, country_code, city, org, asn_num, remote_ip, ))
