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

from io import BytesIO
from json import dumps, loads
from geoip2.database import Reader

from twisted.python.log import msg
from twisted.internet import reactor
from twisted.web import client, http_headers
from twisted.web.client import FileBodyProducer
from twisted.internet.ssl import ClientContextFactory


class WebClientContextFactory(ClientContextFactory):

    def getContext(self, hostname, port):
        return ClientContextFactory.getContext(self)


class QuietHTTP11ClientFactory(client._HTTP11ClientFactory):
    noisy = False


class Output(output.Output):

    def start(self):
        self.host = CONFIG.get('output_nlcvapi', 'host', fallback='localhost')
        self.geoip = CONFIG.getboolean('output_nlcvapi', 'geoip', fallback=True)
        contextFactory = WebClientContextFactory()
        myQuietPool = client.HTTPConnectionPool(reactor)
        myQuietPool._factory = QuietHTTP11ClientFactory
        self.agent = client.Agent(reactor, contextFactory=contextFactory, pool=myQuietPool)

        if self.geoip:
            geoipdb_city_path = CONFIG.get('output_nlcvapi', 'geoip_citydb', fallback='data/GeoLite2-City.mmdb')
            geoipdb_asn_path = CONFIG.get('output_nlcvapi', '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):
        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):
        event['honeypot'] = 'siphoney'
        if self.geoip:
            country, country_code, city, org, asn_num = geolocate(event['src_ip'], self.reader_city, self.reader_asn)
            event['country'] = country
            event['country_code'] = country_code
            event['city'] = city
            event['org'] = org
            event['asn'] = asn_num
        self.postentry(event)

    def postentry(self, entry):
        def cbBody(body):
            return processResult(body)

        def cbPartial(failure):
            """
            Google HTTP Server does not set Content-Length. Twisted marks it as partial
            """
            failure.printTraceback()
            return processResult(failure.value.response)

        def cbResponse(response):
            if response.code in [200, 201]:
                return
            else:
                msg('NLCVAPI response: {}: "{}"'.format(response.code, decode(response.phrase)))
                d = client.readBody(response)
                d.addCallback(cbBody)
                d.addErrback(cbPartial)
                return d

        def cbError(failure):
            failure.printTraceback()

        def processResult(result):
            j = loads(result)
            msg('NLCVAPI response: {}'.format(j['message']))

        headers = http_headers.Headers({
            b'User-Agent': [b'SIP Honeypot'],
            b'Content-Type': [b'application/json']
        })
        d = self.agent.request(
            b'POST',
            self.host.encode('utf-8'),
            headers,
            FileBodyProducer(BytesIO(dumps(entry).encode('utf-8')))
        )
        d.addCallback(cbResponse)
        d.addErrback(cbError)
        return d
