
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 psycopg2 import OperationalError

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


class Output(output.Output):

    def start(self):
        host = CONFIG.get('output_postgres', 'host', fallback='localhost')
        port = CONFIG.getint('output_postgres', 'port', fallback=5432)
        username = CONFIG.get('output_postgres', 'username', fallback='siphoney')
        password = CONFIG.get('output_postgres', 'password')
        database = CONFIG.get('output_postgres', 'database', fallback='siphoney')
        self.debug = CONFIG.getboolean('output_postgres', 'debug', fallback=False)
        self.geoip = CONFIG.getboolean('output_postgres', 'geoip', fallback=True)

        try:
            self.dbh = ConnectionPool(
                'psycopg2',
                database=database,
                user=username,
                password=password,
                host=host,
                port=port,
                cp_min=1,
                cp_max=1
            )
        except OperationalError as e:
            msg('output_postgres: PostgreSQL Error: {}'.format(e))

        if self.geoip:
            geoipdb_city_path = CONFIG.get('output_postgres', 'geoip_citydb', fallback='data/GeoLite2-City.mmdb')
            geoipdb_asn_path = CONFIG.get('output_postgres', '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.dbh.runInteraction(self.connect_event, event)

    def simple_query(self, txn, sql, args, returns_value=True):
        if self.debug:
            msg('output_postgres: postgres query: {} {}'.format(sql, repr(args)))
        result = None
        try:
            txn.execute(sql, args)
            if returns_value:
                result = txn.fetchone()
        except Exception as e:
            msg('output_postgres: postgres Error: {}'.format(e))
        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 = int(r[0])
        else:
            r = self.simple_query(
                txn,
                "INSERT INTO {} ({}) VALUES (%s) RETURNING id".format(table, column),
                (entry, )
            )
            if r:
                id = int(r[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', 'val', value)
            elif key == 'allow':
                allow_id = self.get_id(txn, 'allows', 'val', value)
            elif key == 'call-id':
                callids_id = self.get_id(txn, 'callids', 'val', value)
            elif key == 'contact':
                contact_id = self.get_id(txn, 'contacts', 'val', value)
            elif key == 'content-length':
                conlen_id = self.get_id(txn, 'con_lens', 'val', value)
            elif key == 'content-type':
                contype_id = self.get_id(txn, 'con_types', 'val', value)
            elif key == 'cseq':
                cseq_id = self.get_id(txn, 'cseqs', 'val', value)
            elif key == 'expires':
                expires_id = self.get_id(txn, 'expires', 'val', value)
            elif key == 'from':
                from_id = self.get_id(txn, 'froms', 'val', value)
            elif key == 'max-forwards':
                forward_id = self.get_id(txn, 'forwards', 'val', value)
            elif key == 'supported':
                supported_id = self.get_id(txn, 'supporteds', 'val', value)
            elif key == 'to':
                to_id = self.get_id(txn, 'tos', 'val', value)
            elif key == 'user-agent':
                agent_id = self.get_id(txn, 'agents', 'val', value)
            elif key == 'via':
                via_id = self.get_id(txn, 'vias', 'val', value)
            elif key == 'authorization':
                auth_id = self.get_id(txn, 'authorizations', 'val', value)
            else:
                self.get_id(txn, 'unknown', 'header', key)
        r = self.simple_query(
            txn,
            """
            INSERT INTO requests (accept, allow, callid, contact, con_len, con_type, cseq, expires, frm, forwards, supportd, dest, agent, via, auth)
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            RETURNING id
            """,
            (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, )
        )
        if r:
            id = int(r[0])
        else:
            id = 0
        return id

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

        self.simple_query(
            txn,
            """
            INSERT INTO connections (time_stamp, ip, remote_port, local_host, local_port, request, messg, sensor)
            VALUES (to_timestamp(%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, ),
            False
        )

        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 CONFLICT (ip) DO UPDATE SET
                    country_name = EXCLUDED.country_name,
                    country_iso_code = EXCLUDED.country_iso_code,
                    city_name = EXCLUDED.city_name,
                    org = EXCLUDED.org,
                    org_asn = EXCLUDED.org_asn
                """,
                (remote_ip, country, country_code, city, org, asn_num, ),
                False
            )
