# Copyright (c) 2015 Michel Oosterhof <michel@oosterhof.net>

"""
Splunk HTTP Event Collector (HEC) Connector.
Not ready for production use.
JSON log file is still recommended way to go
"""

from json import dumps, loads
from io import BytesIO

from core import output
from core.config import CONFIG

from zope.interface import implementer

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


@implementer(IPolicyForHTTPS)
class WhitelistContextFactory:
    def creatorForNetloc(self, hostname, port):
        return ssl.CertificateOptions(verify=False)


class Output(output.Output):
    """
    Splunk HEC output
    """

    def start(self):
        self.token = CONFIG.get('output_splunk', 'token')
        self.url = CONFIG.get('output_splunk', 'url').encode('utf-8')
        self.index = CONFIG.get('output_splunk', 'index', fallback='main')
        self.source = CONFIG.get('output_splunk', 'source', fallback='mssqlpot')
        self.sourcetype = CONFIG.get('output_splunk', 'sourcetype', fallback='mssqlpot')
        self.host = CONFIG.get('output_splunk', 'host', fallback=None)
        contextFactory = WhitelistContextFactory()
        self.agent = client.Agent(reactor, contextFactory)

    def stop(self):
        pass

    def write(self, event):
        splunkentry = {}
        if self.index:
            splunkentry['index'] = self.index
        if self.source:
            splunkentry['source'] = self.source
        if self.sourcetype:
            splunkentry['sourcetype'] = self.sourcetype
        if self.host:
            splunkentry['host'] = self.host
        else:
            splunkentry['host'] = event['sensor']
        splunkentry['event'] = event
        self.postentry(splunkentry)

    def postentry(self, entry):
        """
        Send a JSON log entry to Splunk with Twisted
        """
        headers = http_headers.Headers(
            {
                b'User-Agent': [b'MS-SQL Pot Honeypot'],
                b'Authorization': [b'Splunk ' + self.token.encode('utf-8')],
                b'Content-Type': [b'application/json'],
            }
        )
        body = FileBodyProducer(BytesIO(dumps(entry).encode('utf-8')))
        d = self.agent.request(b'POST', self.url, headers, body)

        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 == 200:
                return
            else:
                msg(f'SplunkHEC response: {response.code} {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('SplunkHEC response: {}'.format(j['text']))

        d.addCallback(cbResponse)
        d.addErrback(cbError)
        return d

