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

from geoip2.database import Reader

from psycopg2 import connect

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.conn = connect(
                database=database,
                user=username,
                password=password,
                host=host,
                port=port
            )
        except Exception as e:
            msg(e)

        self.cur = self.conn.cursor()

        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:
                self.reader_city = None
                msg('Failed to open City GeoIP database {}'.format(geoipdb_city_path))

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

    def stop(self):
        self.cur.close()
        self.conn.close()
        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.connect_event(event)

    def simple_query(self, sql, args, returns_value=True):
        if self.debug:
            msg('output_postgres: postgres query: {} {}'.format(sql, repr(args)))
        result = None
        try:
            self.cur.execute(sql, args)
            if returns_value:
                result = self.cur.fetchone()
        except Exception as e:
            msg('output_postgres: postgres Error: {}'.format(e))
        return result

    def get_id(self, table, column, entry):
        r = self.simple_query("SELECT id FROM {} WHERE {} = %s".format(table, column), (entry, ))
        if r:
            id = int(r[0])
        else:
            r = self.simple_query("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, 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('accepts', 'val', value)
            elif key == 'allow':
                allow_id = self.get_id('allows', 'val', value)
            elif key == 'call-id':
                callids_id = self.get_id('callids', 'val', value)
            elif key == 'contact':
                contact_id = self.get_id('contacts', 'val', value)
            elif key == 'content-length':
                conlen_id = self.get_id('con_lens', 'val', value)
            elif key == 'content-type':
                contype_id = self.get_id('con_types', 'val', value)
            elif key == 'cseq':
                cseq_id = self.get_id('cseqs', 'val', value)
            elif key == 'expires':
                expires_id = self.get_id('expires', 'val', value)
            elif key == 'from':
                from_id = self.get_id('froms', 'val', value)
            elif key == 'max-forwards':
                forward_id = self.get_id('forwards', 'val', value)
            elif key == 'supported':
                supported_id = self.get_id('supporteds', 'val', value)
            elif key == 'to':
                to_id = self.get_id('tos', 'val', value)
            elif key == 'user-agent':
                agent_id = self.get_id('agents', 'val', value)
            elif key == 'via':
                via_id = self.get_id('vias', 'val', value)
            elif key == 'authorization':
                auth_id = self.get_id('authorizations', 'val', value)
            else:
                self.get_id('unknown', 'header', key)
        r = self.simple_query(
            """
            INSERT INTO requests (accept, allow, callid, contact, con_len, con_type, cseq, expires, frm, forwards, supportd, dest, agent, via, authorization)
            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, event):
        remote_ip = event['src_ip']
        sensor_id = self.get_id('sensors', 'sname', event['sensor'])
        message_id = self.get_id('messages', 'val', event['message'])
        request_id = self.get_request_id(event['headers'])

        self.simple_query(
            """
            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, )
        )

        if self.geoip:
            country, country_code, city, org, asn_num = geolocate(remote_ip, self.reader_city, self.reader_asn)
            self.simple_query(
                """
                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, country, country_code, city, org, asn_num, ), False)
        self.conn.commit()
