<commit_before>
import urllib3
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

_max_content_bytes = 100000
http_event_collector_debug = False

# Thanks to <NAME> for the http_event_collector class (https://github.com/georgestarcher/)
# Default batch max size to match splunk's default limits for max byte
# See http_input stanza in limits.conf; note in testing I had to limit to 100,000 to avoid http event collector breaking connection
# Auto flush will occur if next event payload will exceed limit
class http_event_collector:

    def __init__(self, token, http_event_server, host='', http_event_port='8088',
                 http_event_server_ssl=True, http_event_collector_ssl_verify=True,
                 max_bytes=_max_content_bytes, proxy=None, timeout=9.05):
        self.timeout = timeout
        self.token = token
        self.batchEvents = []
        self.maxByteLength = max_bytes
        self.currentByteLength = 0
        self.server_uri = []
        self.http_event_collector_ssl_verify = http_event_collector_ssl_verify
        if proxy and http_event_server_ssl:
            self.proxy = {'https': 'https://{0}'.format(proxy)}
        elif proxy:
            self.proxy = {'http': 'http://{0}'.format(proxy)}
        else:
            self.proxy = {}

        # Set host to specified value or default to localhostname if no value provided
        if host:
            self.host = host
        else:
            self.host = socket.gethostname()

        # Build and set server_uri for http event collector
        # Defaults to SSL if flag not passed
        # Defaults to port 8088 if port not passed

        servers = http_event_server
        if not isinstance(servers, list):
            servers = [servers]
        for server in servers:
            if http_event_server_ssl:
                self.server_uri.append(['https://%s:%s/services/collector/event' % (server, http_event_port), True])
            else:
                self.server_uri.append(['http://%s:%s/services/collector/event' % (server, http_event_port), True])

        if http_event_collector_debug:
            print self.token
            print self.server_uri

    def sendEvent(self, payload, eventtime=''):
        # Method to immediately send an event to the http event collector

        headers = {'Authorization': 'Splunk ' + self.token}

        # If eventtime in epoch not passed as optional argument use current system time in epoch
        if not eventtime:
            eventtime = str(int(time.time()))

        # Fill in local hostname if not manually populated
        if 'host' not in payload:
            payload.update({'host': self.host})

        # Update time value on payload if need to use system time
        data = {'time': eventtime}
        data.update(payload)

        # send event to http event collector
        r = requests.post(self.server_uri, data=json.dumps(data), headers=headers,
                          verify=self.http_event_collector_ssl_verify, proxies=self.proxy)

        # Print debug info if flag set
        if http_event_collector_debug:
            logger.debug(r.text)
            logger.debug(data)

    def batchEvent(self, payload, eventtime=''):
        # Method to store the event in a batch to flush later

        # Fill in local hostname if not manually populated
        if 'host' not in payload:
            payload.update({'host': self.host})

        payloadLength = len(json.dumps(payload))

        if (self.currentByteLength + payloadLength) > self.maxByteLength:
            self.flushBatch()
            # Print debug info if flag set
            if http_event_collector_debug:
                print 'auto flushing'
        else:
            self.currentByteLength = self.currentByteLength + payloadLength

        # If eventtime in epoch not passed as optional argument use current system time in epoch
        if not eventtime:
            eventtime = str(int(time.time()))

        # Update time value on payload if need to use system time
        data = {'time': eventtime}
        data.update(payload)

        self.batchEvents.append(json.dumps(data))

    def flushBatch(self):
        # Method to flush the batch list of events

        if len(self.batchEvents) > 0:
            headers = {'Authorization': 'Splunk ' + self.token}
            self.server_uri = [x for x in self.server_uri if x[1] is not False]
            for server in self.server_uri:
                try:
                    r = requests.post(server[0], data=' '.join(self.batchEvents), headers=headers,
                                      verify=self.http_event_collector_ssl_verify,
                                      proxies=self.proxy, timeout=self.timeout)
                    r.raise_for_status()
                    server[1] = True
                    break
                except requests.exceptions.RequestException:
                    #log.info('Request to splunk server "%s" failed. Marking as bad.' % server[0])
                    server[1] = False
                except Exception as e:
                    #log.error('Request to splunk threw an error: {0}'.format(e))
                    pass
            self.batchEvents = []
            self.currentByteLength = 0<commit_msg>DRY #4: unfolded tesseract merge: pull my unifnished HEC re-write back in<commit_after>import socket
import json
import time
import copy
import os

import certifi
import urllib3

# all relating to the persistqueue
import persistqueue
from sqlite3 import DatabaseError as SQLiteDBError
from persistqueue.exceptions import Empty
import shutil
from collections import Counter

import logging
log = logging.getLogger('hubblestack.hec')

_max_content_bytes = 100000
http_event_collector_SSL_verify = True
http_event_collector_debug = False

__version__ = '1.0'

def get_splunk_options(subspace=False, other_defaults=tuple()):
    defaults = other_defaults + (
        ('token', None),
        ('indexer', None),
        ('port', '8088'),
        ('index', None),
        ('custom_fields', []),
        ('sourcetype', 'hubble_fim'),
        ( ('http_event_server_ssl','hec_ssl'), True),
        ('proxy', {}),
        ('timeout', 9.05),
        ('index_extracted_fields', []),
    )

    config_space = __salt__['config.get']('hubblestack:returner:splunk', {})
    if subspace:
        config_space.extend(
            __salt__['config.get']('hubblestack:{0}:returner:splunk'.format(subspace, {})))

    return_opts = {}
    for k,d in defaults:
        if not isinstance(k, (list,tuple)):
            k = [ k ]
        dcls = d.__class__
        v = None
        for cs in config_space:
            for _k in k:
                v = cs.get(_k)
                if v is not None:
                    break
        if v is None:
            v = d
        elif d is not None and v is not None:
            if not isinstance(v, dcls):
                v = dcls(v)
        return_opts[ k[0] ] = v
    return [ return_opts ]

class Payload(object):
    ''' formatters for final payload stringification
        and a convenient place to store retry counter information

        note that formatting a payload is different from formatting an event
        say our event is a line like
          "Sat Aug 18 ProcessName [3828282] DEBUG: oops"
        or perhaps some data like
          {'blah': 'happened"}
        Then the payload would look like this:
          {'host': 'something', 'sourcetype': 'something', 'event': 'Sat Aug 18...'}
          {'host': 'something', 'sourcetype': 'something', 'event': {'blah': 'happened'}}

        For reasons regarding the above, we provide a classmethod to format events:

        p = Payload.format_event({'blah': 'happened'}, sourcetype='blah')
    '''
    host = None

    @classmethod
    def format_event(cls, event, **payload):
        if 'host' not in payload:
            payload['host'] = cls.host
        if 'sourcetype' not in payload:
            payload['sourcetype'] = 'hubble'
        payload['event'] = event
        return cls(payload)

    @classmethod
    def promote(cls, payload, eventtime=''):
        if isinstance(payload, cls):
            return payload
        return Payload(payload, eventtime=eventtime)

    def __init__(self, dat, eventtime=''):
        if self.host is None:
            self.__class__.host = socket.gethostname()

        self.retries = Counter()

        if 'host' not in dat or dat['host'] is None:
            dat['host'] = self.host

        if eventtime:
            dat['time'] = eventtime
        elif 'time' not in dat:
            dat['time'] = str(int(time.time()))

        self.strip_empty_dict_entries(dat)
        self.dat = json.dumps(dat)

    def __repr__(self):
        return 'Payload({0})'.format(self)

    def __str__(self):
        return self.dat

    def __len__(self):
        return len(self.dat)

    @classmethod
    def strip_empty_dict_entries(cls, dat):
        todo = [dat]
        while todo:
            dat = todo.pop(0)
            if isinstance(dat, dict):
                remove_keys = [k for k in dat if not dat[k] and dat[k] not in (0,False) ]
                for k in remove_keys:
                    del dat[k]
                for v in dat.values():
                    if isinstance(v, (list,tuple,dict)):
                        if v not in todo:
                            todo.append(v)
            else:
                for item in dat:
                    if isinstance(item, (list,tuple,dict)):
                        if item not in todo:
                            todo.append(item)

class DiskQueue(persistqueue.SQLiteQueue):
    max_size = 10000

    def __init__(self, fname='/var/cache/hubble/pq.db'):
        self._last_rfq = 0
        try:
            super(DiskQueue, self).__init__(fname)
        except SQLiteDBError:
            # XXX: is this dangerous?
            if os.path.isdir(fname):
                shutil.rmtree(fname)
            # one retry
            super(HEC, self).__init__(fname)

    def put(self, item):
        # XXX: max_size is in items, not bytes
        # is this better in bytes? is the limit too high? too low?
        if self.size < self.max_size:
            super(DiskQueue, self).put(item)
            self._last_rfq = 0
        else:
            t = time.time()
            dt = t - self._last_rfq
            if dt > queue_overflow_msg_delay:
                log.error("refusing to queue event item: queue full")
                self._last_rfq = t

    def next(self):
        return self.get(block=False)

    def __iter__(self):
        while True:
            try:
                yield self.next()
            except Empty:
                break

# Thanks to <NAME> for the http_event_collector class (https://github.com/georgestarcher/)
# Default batch max size to match splunk's default limits for max byte
# See http_input stanza in limits.conf; note in testing I had to limit to
# 100,000 to avoid http event collector breaking connection Auto flush will
# occur if next event payload will exceed limit

class HEC(object):

    def __init__(self, token, http_event_server, host='', http_event_port='8088',
        http_event_server_ssl=True, max_bytes=_max_content_bytes, proxy=None, timeout=9.05,
        log_http_exceptions=False, log_other_exceptions=False):

        self.payload_retry_before_fail =  5
        self.queue_overflow_msg_delay = 10
        self.retry_diskqueue_interval = 60
        self._last_dq_flush_attempt = 0

        self.timeout = timeout
        self.token = token
        self.batchEvents = []
        self.maxByteLength = max_bytes
        self.currentByteLength = 0
        self.server_uri = []
        self.log_http_exceptions = log_http_exceptions
        self.log_other_exceptions = log_other_exceptions

        if proxy and http_event_server_ssl:
            self.proxy = {'https': 'https://{0}'.format(proxy)}
        elif proxy:
            self.proxy = {'http': 'http://{0}'.format(proxy)}
        else:
            self.proxy = {}

        # Set host to specified value or default to localhostname if no value provided
        if host:
            self.host = host
        else:
            self.host = socket.gethostname()

        Payload.host = self.host

        # Build and set server_uri for http event collector
        # Defaults to SSL if flag not passed
        # Defaults to port 8088 if port not passed

        servers = http_event_server
        if not isinstance(servers, list):
            servers = [servers]
        for server in servers:
            if http_event_server_ssl:
                self.server_uri.append(['https://%s:%s/services/collector/event' % (server, http_event_port), True])
            else:
                self.server_uri.append(['http://%s:%s/services/collector/event' % (server, http_event_port), True])

        if http_event_collector_debug:
            print self.token
            print self.server_uri

        # build headers once
        self.headers = urllib3.make_headers( keep_alive=True,
            user_agent='hubble-hec/{0}'.format(__version__),
            accept_encoding=True)
        self.headers.update({ 'Content-Type': 'application/json',
            'Authorization': 'Splunk {0}'.format(self.token) })

        # retries can be made much more flexible than shown here theoretically,
        # it could take the load off overloaded servers through the backoff and
        # improve overall throughput at those (usually transient) bottlenecks
        # -- the number 3 was chosen essentially at random
        pm_kw = {
            'timeout': self.timeout,
            'retries': urllib3.util.retry.Retry(
                total=5, redirect=10, backoff_factor=3,
                connect=self.timeout, read=self.timeout,
                respect_retry_after_header=True)
        }

        if http_event_collector_SSL_verify:
            pm_kw.update({'cert_reqs': 'CERT_REQUIRED', 'ca_certs': certifi.where()})
        else:
            urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

        if self.proxy:
            self.pool_manager = urllib3.ProxyManager(self.proxy, **pm_kw)
        else:
            self.pool_manager = urllib3.PoolManager(**pm_kw)

        self.queue = DiskQueue()

    def queueEvent(self, dat, eventtime=''):
        if not isinstance(dat, Payload):
            payload = Payload(dat, eventtime)
        self.queue.put(payload)

    def flushQueue(self):
        for i in self.queue:
            self.batchEvent(i)
        self.flushBatch()

    def _attempt_flush_queue(self):
        t = time.time()
        dt = t - self._last_dq_flush_attempt
        self._last_dq_flush_attempt = time.time()
        if dt >= self.retry_diskqueue_interval and self.queue.size:
            log.debug('flushing queue sz=%d', self.queue.size)
            self.flushQueue()

    def _requeue(self, payloads, server):
        server_maybe_bad = False
        for pload in payloads:
            pload.retries[ server ] += 1
            if pload.retries[ server ] < self.payload_retry_before_fail:
                self.queue.put(pload)
            else:
                server_maybe_bad = True
        return server_maybe_bad

    def _send(self, *payload):
        data = ' '.join([ str(x) for x in payload ])
        self.server_uri = [x for x in self.server_uri if x[1] is not False]
        for server in self.server_uri:
            try:
                r = self.pool_manager.request('POST', server[0], body=data, headers=self.headers)
            except urllib3.exceptions.LocationParseError as e:
                if self.log_other_exceptions:
                    log.error('server uri parse error "{0}": {1}'.format(server[0], e))
                server[1] = False
                continue
            except Exception as e:
                if self.log_other_exceptions:
                    log.error('Request to splunk threw an error: {0}'.format(e))
                server[1] = False
                continue
            if r.status >= 400:
                server_maybe_bad = self._requeue(payload, server[0])
                if server_maybe_bad:
                    if self.log_http_exceptions:
                        log.info('Request to splunk server "%s" failed. Marking as bad.' % server[0])
                    server[1] = False
                #self._requeue(self, payload, 'blah')
            else:
                return r

    def sendEvent(self, payload, eventtime=''):
        self._attempt_flush_queue()
        r = self._send( Payload.promote(payload, eventtime=eventtime) )

        if http_event_collector_debug:
            if r is not None and hasattr(r, 'text'):
                log.debug(r.text)
            log.debug(payload)


    def batchEvent(self, dat, eventtime=''):
        self._attempt_flush_queue()
        payload = Payload.promote(dat, eventtime)

        if (self.currentByteLength + len(payload)) > self.maxByteLength:
            self.flushBatch()
            if http_event_collector_debug:
                log.debug('auto flushing')
        else:
            self.currentByteLength = self.currentByteLength + len(payload)
        self.batchEvents.append(payload)


    def flushBatch(self):
        if self.batchEvents:
            r = self._send( *self.batchEvents )
            self.batchEvents = []
            self.currentByteLength = 0

http_event_collector = HEC
{'fraction_non_alphanumeric': 0.09433962264150944, 'fraction_numerical': 0.007708237459733088, 'mean_word_length': 2.035090782122905, 'pattern_counts': {'":': 1, '<': 7, '<?xml version=': 0, '>': 11, 'https://': 6, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3579299', 'n_tokens_mistral': 5032, 'n_tokens_neox': 4816, 'n_words': 1354}