#!/usr/bin/python3
"""
Author: SYNchroACK <synchroack@protonmail.ch>
"""

import argparse
import json
import sys

try:
    from elasticsearch import Elasticsearch
except ImportError:
    print("[-] Please install Elasticsearch using the following command: 'pip install elasticsearch'.", file=sys.stderr)
    sys.exit(-1)


TYPES = {
    "Accuracy": "long",
    "ASN": "integer",
    "Base64": "text",
    "Boolean": "boolean",
    "ClassificationType": "keyword",
    "DateTime": "date",
    "Float": "long",
    "FQDN": "text",
    "Integer": "integer",
    "IPAddress": "ip",
    "IPNetwork": "text",
    "JSON": "text",
    "JSONDict": "nested",
    "LowercaseString": "text",
    "Registry": "text",
    "String": "text",
    "UppercaseString": "text",
    "TLP": "text",
    "URL": "text"
}


def __mapping_properties_from_harmonization(properties):
    # perform your customization on mapping (if needed)
    del properties['extra']
    return properties


def mapping_properties_from_harmonization(harmonization, replacement_char):
    properties = dict()
    err = None

    for field, field_options in harmonization['event'].items():

        field = field.replace('.', replacement_char)

        if field_options['type'] in TYPES.keys():
            field_type = TYPES[field_options['type']]
        else:
            field_type = "text"  # fallback type
            err = -1

        properties[field] = dict(type=field_type)

    return __mapping_properties_from_harmonization(properties), err


def create_mapping(harmonization, replacement_char):

    config = {"enabled": False}

    properties, err = mapping_properties_from_harmonization(harmonization, replacement_char)

    data = {
        "mappings": {
            "properties": properties
        }
    }

    return data, err


def send_mapping(host, index, data):
    err = None
    response = None

    try:
        es = Elasticsearch([host], verify_certs=True)
        response = es.indices.create(index=index, ignore=400, body=data)
    except Exception:
        err = -1

    return response, err


def create_template(mapping, patterns=None):
    """
    Constructs an Elasticsearch index template from a given mapping.
    :param mapping: The mapping to use in the template.
    :param patterns: (Optional) A list of patterns for this template to match. Default: ["intelmq-*"]
    :return: A dict containing the new index template.
    """

    patterns = patterns or ["intelmq-*"]

    if len(patterns) > 1:
        print("[!] Notice: The template contains multiple index patterns ({}). "
              "Multiple index patterns are only supported on Elasticsearch >= 6.0.0."
              "If your Elasticsearch instance does not support multiple index patterns, "
              "only the first value ({}) will be used.".format(patterns, patterns[0]))

    template = {
        "index_patterns": patterns,
        "template": patterns[0],  # For backward compatibility with ES < 6.0.0. Ignored if ES >= 6.0.0
        # "settings": {
        #     "number_of_shards": 5,
        #     ...
        # },
        "mappings": mapping.get("mappings")
    }

    return template


def send_template(host, name, template):
    """
    Send a mapping template to Elasticsearch
    :param host: The Elasticsearch host
    :param name: The name to use for the template
    :param template: The template (as a dict) to send
    :return: response, err as returned by the Elasticsearch client.
    """
    err = None
    response = None

    try:
        es = Elasticsearch([host], verify_certs=True)
        response = es.indices.put_template(name=name, body=template)
    except:
        err = -1

    return response, err


if __name__ == "__main__":

    parser = argparse.ArgumentParser(
        description='Elastic Mapper tool',
    )

    parser.add_argument('--harmonization-file',
                        action="store",
                        dest="harmonization_file",
                        metavar="<filepath>",
                        required=True,
                        help='harmonization file')

    parser.add_argument('--harmonization-fallback',
                        action="store_true",
                        dest="harmonization_fallback",
                        required=False,
                        default=False,
                        help='harmonization fallback to `text` type')

    parser.add_argument('--host',
                        action="store",
                        dest="host",
                        metavar="<ip>",
                        required=False,
                        help='elasticsearch server IP')

    parser.add_argument('--index',
                        action="store",
                        dest="index",
                        default="intelmq",
                        required=False,
                        help='elasticsearch index name, or template name if using a template')

    parser.add_argument('--index-template',
                        action="store_true",
                        dest="index_template",
                        required=False,
                        default=False,
                        help='whether to store the mapping as a template for new indices')

    parser.add_argument('--replacement-char',
                        action="store",
                        dest="replacement_char",
                        default="_",
                        required=False,
                        help='replacement char for all fields (replace . with _ by default)')

    parser.add_argument('--output',
                        action="store",
                        dest="output",
                        metavar="<filepath>",
                        required=False,
                        help='write a copy of applied mapping or template to file')

    arguments = parser.parse_args()

    with open(arguments.harmonization_file) as fp:
        harmonization = json.load(fp)

    data, err = create_mapping(harmonization, arguments.replacement_char)

    if err:
        if arguments.harmonization_fallback:
            print("[-] Found a field type not recognizable by elasticmapper. Fallback to `text` type.", file=sys.stderr)
        else:
            print("[-] Found a field type not recognizable by elasticmapper.", file=sys.stderr)
            print("[-] Please fix elasticmapper or specify '--harmonization-fallback' parameter to use the fallback type `text`.", file=sys.stderr)
            print("[-] Supported field types by elasticmapper: %s" % ", ".join(TYPES.keys()), file=sys.stderr)
            sys.exit(-2)

    # Optionally construct an index template from the generated mapping
    if arguments.index_template:
        data = create_template(data)

    if arguments.output:
        with open(arguments.output, 'w') as fp:
            fp.write(json.dumps(data, indent=4))

        object_type = "Template" if arguments.index_template else "Mapping"
        print("[+] {} has been written to output file.".format(object_type))

    if arguments.host:

        if arguments.index_template:
            # Optionally, send the mapping as a template for new indices
            response, err = send_template(arguments.host, arguments.index, data)
        else:
            # Default - create a mapping for only one named index
            response, err = send_mapping(arguments.host, arguments.index, data)

        if err:
            print("\n[-] Could not send mapping to Elasticsearch.", file=sys.stderr)
            print("[-] Please check if host parameter is correct or if host is allowing connections to Elasticsearch REST API.", file=sys.stderr)
            sys.exit(-3)
        else:
            print("[+] Mapping has been sent to Elasticsearch.")
            print("[+] Elasticsearch Response:\n")
            print(response)
            print("")
