
from __future__ import absolute_import

from sys import version_info

from core import output
from core.config import CONFIG

from twisted.python.log import msg

PY2 = version_info[0] < 3

if PY2:
    import elasticsearch    # type: ignore
    from elasticsearch import Elasticsearch, NotFoundError  # type: ignore
    version_module = elasticsearch
    version_attr = '__version__'
else:
    import elasticsearch8
    from elasticsearch8 import Elasticsearch, NotFoundError
    version_module = elasticsearch8
    version_attr = '__versionstr__'


class Output(output.Output):

    def start(self):
        host = CONFIG.get('output_elastic', 'host', fallback='localhost')
        port = CONFIG.getint('output_elastic', 'port', fallback=9200)
        username = CONFIG.get('output_elastic', 'username', fallback=None)
        password = CONFIG.get('output_elastic', 'password', fallback=None)
        use_ssl = CONFIG.getboolean('output_elastic', 'ssl', fallback=False)
        verify_certs = CONFIG.getboolean('output_elastic', 'verify_certs', fallback=False)
        ca_certs = CONFIG.get('output_elastic', 'ca_certs', fallback=None)
        self.index = CONFIG.get('output_elastic', 'index', fallback='siphoney')
        self.pipeline = CONFIG.get('output_elastic', 'pipeline', fallback='geoip')

        scheme = 'https' if use_ssl else 'http'
        hosts = [{'host': host, 'port': port, 'scheme': scheme}]
        kwargs = {}
        if username and password:
            kwargs['http_auth' if PY2 else 'basic_auth'] = (username, password)
        if use_ssl:
            kwargs['verify_certs'] = verify_certs
            if verify_certs and ca_certs:
                kwargs['ca_certs'] = ca_certs

        # Create client
        self.es = Elasticsearch(hosts=hosts, **kwargs)

        # Detect major version from module attribute
        ver_attr_value = getattr(version_module, version_attr)
        if isinstance(ver_attr_value, tuple):
            # Python 2.7 elasticsearch 7.x
            self.es_ver = ver_attr_value[0]
        else:
            # Python 3.6+ elasticsearch8
            self.es_ver = int(str(ver_attr_value).split('.')[0])

        # Test connection
        try:
            info = self.es.info()
            msg('output_elastic: connected to ES version {}'.format(info['version']['number']))
        except Exception as e:
            msg('output_elastic: connection failed: {}'.format(repr(e)))
            raise

        # Ensure index exists
        if not self.es.indices.exists(index=self.index):
            self.es.indices.create(index=self.index)

        # Setup geoip
        if self.pipeline == 'geoip':
            self._setup_geoip_mapping()
            self._setup_geoip_pipeline()

    def _setup_geoip_mapping(self):
        try:
            body = {
                'properties': {
                    'geo': {
                        'properties': {
                            'location': {
                                'type': 'geo_point'
                            }
                        }
                    }
                }
            }
            if self.es_ver >= 8:
                self.es.indices.put_mapping(index=self.index, properties=body['properties'])
            else:
                self.es.indices.put_mapping(index=self.index, body=body)
        except Exception as e:
            msg('output_elastic: mapping setup failed: {}'.format(repr(e)))

    def _setup_geoip_pipeline(self):
        try:
            self.es.ingest.get_pipeline(id=self.pipeline)
        except NotFoundError:
            body = {
                'description': 'Add geoip info',
                'processors': [
                    {
                        'geoip': {
                            'field': 'src_ip',
                            'target_field': 'geo',
                            'database_file': 'GeoLite2-City.mmdb'
                        }
                    }
                ]
            }
            try:
                if self.es_ver >= 8:
                    self.es.ingest.put_pipeline(
                        id=self.pipeline,
                        processors=body['processors'],
                        description=body['description']
                    )
                else:
                    self.es.ingest.put_pipeline(id=self.pipeline, body=body)
            except Exception as e:
                msg('output_elastic: pipeline creation failed: {}'.format(repr(e)))

    def stop(self):
        pass

    def write(self, entry):
        try:
            if self.es_ver >= 8:
                self.es.index(index=self.index, document=entry, pipeline=self.pipeline)
            else:
                self.es.index(index=self.index, body=entry, pipeline=self.pipeline)
        except Exception as e:
            msg('output_elastic: write failed: {}'.format(repr(e)))
            if hasattr(e, 'info'):
                msg('output_elastic info: {}'.format(e.info))
