"""
Send SSH logins to SANS DShield.
See https://isc.sans.edu/ssh.html
"""

from base64 import b64decode, b64encode
from hashlib import sha1, sha256, md5
from hmac import new
from re import compile
from time import strftime

from dateutil.parser import parse
from requests import request

from core import output
from core.config import CONFIG

from twisted.internet import reactor
from twisted.internet import threads
from twisted.python.log import msg


class Output(output.Output):
    """
    dshield output
    """

    debug: bool = False
    userid: str
    batch_size: int
    batch: list

    def start(self):
        self.auth_key = CONFIG.get('output_dshield', 'auth_key')
        self.userid = CONFIG.get('output_dshield', 'userid')
        self.batch_size = CONFIG.getint('output_dshield', 'batch_size')
        self.debug = CONFIG.getboolean('output_dshield', 'debug', fallback=False)
        self.batch = []  # This is used to store login attempts in batches

    def stop(self):
        pass

    def write(self, event):
        if event['operation'].lower() == 'login':
            date = parse(event['timestamp'])
            self.batch.append(
                {
                    'date': str(date.date()),
                    'time': date.time().strftime('%H:%M:%S'),
                    'timezone': strftime('%z'),
                    'source_ip': event['src_ip'],
                    'user': event['username'],
                    'password': event['password'],
                }
            )

            if len(self.batch) >= self.batch_size:
                batch_to_send = self.batch
                self.submit_entries(batch_to_send)
                self.batch = []

    def transmission_error(self, batch):
        self.batch.extend(batch)
        if len(self.batch) > self.batch_size * 2:
            self.batch = self.batch[-self.batch_size :]

    def submit_entries(self, batch):
        """
        Large parts of this method are adapted from kippo-pyshield by jkakavas
        Many thanks to their efforts. https://github.com/jkakavas/kippo-pyshield
        """
        # The nonce is predefined as explained in the original script :
        # trying to avoid sending the authentication key in the "clear" but
        # not wanting to deal with a full digest like exchange. Using a
        # fixed nonce to mix up the limited userid.
        _nonceb64 = 'ElWO1arph+Jifqme6eXD8Uj+QTAmijAWxX1msbJzXDM='

        log_output = ''
        for attempt in self.batch:
            log_output += '{}\t{}\t{}\t{}\t{}\t{}\n'.format(
                attempt['date'],
                attempt['time'],
                attempt['timezone'],
                attempt['source_ip'],
                attempt['user'],
                attempt['password'],
            )

        nonce = b64decode(_nonceb64)
        digest = b64encode(
            new(
                nonce + self.userid.encode('ascii'),
                b64decode(self.auth_key),
                sha256,
            ).digest()
        )
        auth_header = 'credentials={} nonce={} userid={}'.format(
            digest.decode('ascii'), _nonceb64, self.userid
        )
        headers = {'X-ISC-Authorization': auth_header, 'Content-Type': 'text/plain'}

        if self.debug:
            msg('dshield: posting: {}'.format(headers))
            msg('dshield: posting: {}'.format(log_output))

        req = threads.deferToThread(
            request,
            method='PUT',
            url='https://secure.dshield.org/api/file/sshlog',
            headers=headers,
            timeout=10,
            data=log_output,
        )

        def check_response(resp):
            failed = False
            response = resp.content.decode('utf8')

            if self.debug:
                msg('dshield: status code {}'.format(resp.status_code))
                msg('dshield: response {}'.format(resp.content))

            if resp.ok:
                sha1_regex = compile(r'<sha1checksum>([^<]+)<\/sha1checksum>')
                sha1_match = sha1_regex.search(response)
                sha1_local = sha1()
                sha1_local.update(log_output.encode('utf8'))
                if sha1_match is None:
                    msg(
                        'dshield: ERROR: Could not find sha1checksum in response: {}'.format(response)
                    )
                    failed = True
                elif sha1_match.group(1) != sha1_local.hexdigest():
                    msg('dshield: ERROR: SHA1 Mismatch {} {} .'.format(sha1_match.group(1), sha1_local.hexdigest()))
                    failed = True

                md5_regex = compile(r'<md5checksum>([^<]+)<\/md5checksum>')
                md5_match = md5_regex.search(response)
                md5_local = md5()
                md5_local.update(log_output.encode('utf8'))
                if md5_match is None:
                    msg('dshield: ERROR: Could not find md5checksum in response')
                    failed = True
                elif md5_match.group(1) != md5_local.hexdigest():
                    msg('dshield: ERROR: MD5 Mismatch {} {} .'.format(md5_match.group(1), md5_local.hexdigest()))
                    failed = True

                msg('dshield: SUCCESS: Sent {} bytes worth of data to secure.dshield.org'.format(log_output))
            else:
                msg('dshield ERROR: error {}.'.format(resp.status_code))
                msg('dshield response was {}'.format(response))
                failed = True

            if failed:
                # Something went wrong, we need to add them to batch.
                reactor.callFromThread(self.transmission_error, batch)

        req.addCallback(check_response)
