#!/bin/env python3
#
# Helper to create cloudprober's config from simple YAML with the following structure
#
# blackbox:
#   targets:
#     host.mask: *dns*  # all dns servers
#       # we have to use host.match: regex since powerdns supports only '*'
#       # to replace several characters. And it's the way to filter hosts more concrete
#     host.match: .*-prd-dns
#   {any.name}:
#     labels:
#       {key}: {value}
#     {more.specific.name.wich.inherites.labels}:
#       labels:
#         {key.in.this.scope}: {value}
#       services:
#         - name: {any.service.name}
#           rule: {schema}://*-env-svc*.*.example.net | status:200-299 | content:"Some string"
#           {svc.label.key}: {value}
#           {another.label.key}: {another_value}
#
#
# Integrated with PowerDNS. Generates list of hosts specified by mask.
#

import os
import sys
import yaml
import argparse
import json
from urllib.parse import urlparse
from difflib import unified_diff
# import axconfig
import re
from collections import OrderedDict
from copy import deepcopy
from contextlib import contextmanager


INHERITED_KEYS = ['config', 'labels', 'targets', 'proto', 'alerts', 'teams']
OVERRIDE_KEYS = []
# OVERRIDE_KEYS = ['config']


surfacers = [{
    "type": "PROMETHEUS",
    "ignore_metrics_with_label": [{
      "key": "probe",
      "value": "sysvars"
    }]
  }]


labels_mandatory = (
    "type",          # internal, external, core
    "product",       # replibit, cloudfinder, shadowprotect, anchor, brc, bff, veeam, appassure
    "role",          # balancer, master, web, backend, vault, etc...
    "env",           # dev, qa, stg, prod
    "severity",      # info, warning, high, critical
    "datacenter"
)


def merge(a: dict, b: dict) -> dict:
    tmp = deepcopy(a)
    for key in b:
        if key in tmp:
            if isinstance(tmp[key], dict) and isinstance(b[key], dict):
                tmp[key] = merge(tmp[key], b[key])
            else:
                tmp[key] = deepcopy(b[key])
        else:
            tmp[key] = deepcopy(b[key])
    return tmp


@contextmanager
def context(ctx={}, tree=None):
    try:
        tmp = deepcopy(ctx)
        if isinstance(tree, dict):
            for key in ctx['OVERRIDE_KEYS']:
                if key in tree:
                    tmp[key] = deepcopy(tree[key])
            for key in ctx['INHERITED_KEYS']:
                if key in tree:
                    if key in tmp:
                        tmp[key] = merge(tmp[key], tree[key])
                    else:
                        tmp[key] = deepcopy(tree[key])
        yield tmp
    finally:
        pass


def dict2kvlist(d: dict, keyfield: str = "key") -> list:
    if isinstance(d, dict):
        return [{keyfield: k, "value": v} for k, v in d.items()]
    return d


def kvlist2dict(lst: list, keyfield: str = "key") -> dict:
    if isinstance(lst, list):
        return {i[keyfield]: i["value"] for i in lst}
    return lst


# Adding labels into probe config
#
#   probe {
#     ...
# +   additional_label {
# +     key:   {key}
# +     value: {val}
# +   }
#   }
#
def render_labels(labels: dict):
    if 'datacenter' not in labels:
        labels["datacenter"] = "@target.label.datacenter@"

    return { "additional_label": dict2kvlist(labels) }


#
#  Generates http params
#
def http_params(url):
    default_port = 443 if url.scheme == "https" else 80
    return {
               "http_probe": {
                   "protocol": url.scheme.upper(),
                   "relative_url": url.path,
                   "port": url.port if url.port is not None else default_port
               }
           }


#
#  Generates tcp params
#
def tcp_params(url):
    return {
               "tcp_probe": {
                   "port": url.port
               }
           }


#
#  Generates dns params
#
def dns_params(url):
    return {
               "dns_probe": {
                   "resolved_domain": url.path.lstrip('/')
               }
           }


#
# Generates params for builting extension protocols
#
def extension_params(url, proto):
    opts = {
        "port": int(url.port),
    }
    prober = url.scheme
    x = {**proto, **opts}
    return {
               f"[axcprober.{prober}_probe]": x
           }


#
# Generates params for external modules
#
def external_params(url, proto):
    opts = {
        "proto": url.scheme,
        "target": "@target@",
    }
    if url.port is not None:
        opts["port"] = str(url.port)
    x = {**proto}
    x["mode"] = x["mode"].upper()
    x["options"] = dict2kvlist({**opts, **x["options"]}, keyfield="name")
    return {
        "external_probe": x
    }

# Save all the "... | FILTER | FILTER ..." as an OrderedDict
# Example:
#   ... | status:200-299 | content:"String"
#   will be represented as
#   {
#     "status":  "200-299"
#     "content": "String"
#   }
def parse_filters(filters_list):
    r = OrderedDict()
    for filter in filters_list:
        filter = filter.lstrip(' \t').rstrip(' \t')
        k, v = filter.split(':')
        r[k] = v
    return r


# Gets datacenter from an FQDN to use as the `datacenter` label later
def datacenter(fqdn: str):
    m = re.match(r'^.*\.([^\.]+)\.efscloud\.net$', fqdn)
    if m:
        return m.group(1)
    return None


def http_validators(filters):
    i = 1
    ret = dict()
    if filters:
        if "status" in filters:
            ret = {
                "validator": [
                    {
                        "name": f"status_{i}",
                        "http_validator": {
                            "success_status_codes": filters["status"]
                        }
                    }
                ]
            }
        else:
            ret = {
                "validator": [
                    {
                        "name": f"status_{i}",
                        "http_validator": {
                            "success_status_codes": "200-299"
                        }
                    }
                ]
            }
        if "content" in filters:
            r = {
                "name": f"content_{i}",
                "integrity_validator": {
                    "pattern_string": filters["content"].strip('"')
                }
            }
            ret["validator"].append(r)
    else:
        ret = {
            "validator": [
                {
                    "name": f"status_{i}",
                    "http_validator": {
                        "success_status_codes": "200-299"
                    }
                }
            ]
        }
    return ret


#
#  Parse rule
#
def parse_rule(ctx, rule: str, name: str):  # -> ???
    targets = ctx["targets"] if "targets" in ctx else {}
    proto = ctx["proto"] if "proto" in ctx else {}
    r = None
    # Since rule is the pipeline of "URL | FILTER | FILTER ..."
    # split it by ' | ' (whitespaces is necessary)
    pipeline = rule.split(' | ')

    # URL is the 0 element of the pipeline
    # It will throw IndexError in case nothing in the rule
    # Also remove all the whitespaces/tabs around the url
    url = urlparse(pipeline[0].rstrip(' \t').lstrip(' \t'))

    filters = OrderedDict()
#    TODO: add validators
    # Parse filters as an OrderedDict if we have them
    if len(pipeline) > 1:
        filters = parse_filters(pipeline[1:])
#        render_validators(pipeline[1:])

    type = "HTTP"

    if url.scheme == "https":
        # type HTTP is used for both http:// and https://
        # in the second case proto = HTTPS will be added in http_params
        type = "HTTP"
    elif url.scheme in proto:
        if "mode" in proto[url.scheme]:
            type = "EXTERNAL"
        else:
            type = "EXTENSION"
    else:
        type = url.scheme.upper()

    m = re.match(r'^\$\{(.+)\}$', url.hostname)
    if m:
        # If hostname is a ${variable} lookup service

        r = {
            "type": type,
            "targets": {
                "[axcprober.service]": {
                    "service": f"{name}"
                }
            }
        }

    # target hostname isn't variable
    else:
        # host_search_file(url.hostname, filters, f'{name}.json')

        r = {
            "type": type,
            "targets": {
                "host_names": f"{url.hostname}"
            }
        }

    v = {}

    p = {}

    if url.scheme in ["http", "https"]:
        v = http_validators(filters)
        p = http_params(url)

    if url.scheme in ["tcp"]:
        p = tcp_params(url)

    if url.scheme in ["dns"]:
        p = dns_params(url)

    if url.scheme in proto:
        if "mode" in proto[url.scheme]:
            p = external_params(url, proto[url.scheme])
        else:
            p = extension_params(url, proto[url.scheme])

    return {**r, **p, **v}


def render_alerts(alerts: dict, teamcfg: dict, severity: str):
    import copy
    __providers = copy.deepcopy(alerts['providers'])
    __rules = alerts['rules']
    __config = alerts['config'] if 'config' in alerts else {}
    __severity = "UNKNOWN_SEVERITY"
    if "high" == severity:
        __severity = "ERROR"
    elif severity in ["critical", "error", "warning", "info"]:
        __severity = severity.upper()
    r = list()
    for rule in __rules:
        if "condition" in rule and 2 == len(rule["condition"].split('/')):
            failures, total = rule["condition"].split('/')
            p = "unknown"
            # Cloudprober supports closing alerts for PagerDuty and Opsgenie
            # so if we use several alerts it will call closing alert function
            # for any other type of notification and as a result won't close
            # PagerDuty or Opsgenie alert because it will be deleted from
            # current alerts list while cloudprober processes other type
            if "default" in rule['notify']:
                p = rule['notify']["default"]
            if severity in rule['notify']:
                p = rule['notify'][severity]
#                for x in rule['notify'][severity]:
#                    if x not in p2send:
#                        p2send.append(x)
#            __local_config = rule['config'] if 'config' in rule else {}
#            __local_config = {**__config, **__local_config}
            #for p in p2send: # uncomment this when cloudprober will be ready
                              # to work with different alert types
            provider = merge(__providers[p], teamcfg[p] if p in teamcfg else {})
            __local_config = {**__config, **provider.pop("config", {})}
            tmp = {
                "condition": {
                    "failures": int(failures),
                    "total": int(total)
                },
                "severity": __severity,
                "notify": {
                    p: dict(provider)
                },
            }
            for k in __local_config:
                tmp[k] = __local_config[k]
            r.append(tmp)
            # end of 'for p in p2send' loop

    return {"alert": r}


#
# Render "probe { ... }" for the service
#
def render_svcs(ctx, svcs: list) -> str:
    r = list()
    for svc in svcs:
        with context(ctx, svc) as _ctx:
            __labels = _ctx["labels"] if "labels" in _ctx else {}
            __config = _ctx["config"] if "config" in _ctx else {}
            __alerts = _ctx["alerts"] if "alerts" in _ctx else {}
            __teams = _ctx["teams"] if "teams" in _ctx else {}
            name = None
            if 'name' in svc:
                name = svc['name']
            if 'rule' in svc:
                rule_content = parse_rule(_ctx, svc['rule'], svc['name'])

            rlabels = render_labels(__labels)
            ralerts = {}
            __teamcfg = __teams[__labels["team"]] if "team" in __labels and __labels["team"] in __teams else {}
            if "severity" in __labels:
                ralerts = render_alerts(__alerts, __teamcfg, __labels["severity"], )

            tmp = {
                "name": name,
                "interval": "60s",
                "timeout": "3s"
            }
            r.append({**tmp, **__config, **rule_content, **rlabels, **ralerts})
    return r


def proto_fmt(proto: dict) -> dict:
    __proto = {**proto}
    for i in __proto:
        __proto[i]["mode"] = __proto[i]["mode"].upper()
    return __proto


def render(ctx, tree: dict) -> str:
    r = list()

    with context(ctx, tree) as _ctx:
        for entry in tree:
            if entry in INHERITED_KEYS or entry in OVERRIDE_KEYS:
                pass

            elif 'services' == entry:
                r += render_svcs(_ctx, tree['services'])

            else:
                r += render(_ctx, tree[entry])
    return r


def to_textpb(obj, level=0, indent=2, upkey=""):
    lines = list()
    if isinstance(obj, dict):
        for k in obj:
            if isinstance(obj[k], list):
                for v in obj[k]:
                    lines.append(' '*indent*level + f'{k}' + ' {')
                    lines.extend(to_textpb(v, level=level+1, upkey=k))
                    lines.append(' '*indent*level + '}')
            elif isinstance(obj[k], dict):
                lines.append(' '*indent*level + f'{k}' + ' {')
                lines.extend(to_textpb(obj[k], level=level+1, upkey=k))
                lines.append(' '*indent*level + '}')
            elif isinstance(obj[k], str) and re.match(r'^[A-Z]+$', obj[k]) is None:
                s = obj[k]
                s = s.replace('"','\\"')
                s = s.replace('\n','\\n')
                lines.append(' '*indent*level + f'{k}: "{s}"')
            else:
                lines.append(' '*indent*level + f'{k}: {obj[k]}')
    return lines


def parse_args():
    parser = argparse.ArgumentParser( prog='axcprober-cfggen',
                                      description='Generate cloudprober config')
    parser.add_argument('-o', '--out', nargs='?', type=str, help="File to output config to")
    return parser.parse_args()


if __name__ == "__main__":

    args = parse_args()
    stdinbuf = ""
    cfg = None

    for line in sys.stdin:
        stdinbuf += line

    try:
        cfg = yaml.safe_load(stdinbuf)
        if args.out:
            current = ""
            if os.path.isfile(args.out):
                with open(args.out, 'r') as f:
                    current = f.read().splitlines()

            content = to_textpb({
                "probe": render(ctx={
                    "INHERITED_KEYS": INHERITED_KEYS,
                    "OVERRIDE_KEYS": OVERRIDE_KEYS,
                    }, tree=cfg),
                "surfacer": surfacers
            })

            diff = list()
            for line in unified_diff(current, content):
                diff.append(line)
            if len(diff) > 0:
                saltret = {
                        'changed': True,
                        'comment': '\n'.join(diff)
                        }
                print(json.dumps(saltret))
            with open(args.out, 'w') as f:
                f.write('\n'.join(content))
        else:
            # cloudprober supports different types of config files:
            #   - textpb
            #   - yaml
            #   - json
            # but it has issues to marshalling yaml and json formats to its'
            # native textpb. For example http_notify config won't be converted
            # correctly because of escaping symbols. So at this moment
            # we shouldn't use anything, excepting native textpb format
            print('\n'.join(to_textpb({
                "probe": render(ctx={
                    "INHERITED_KEYS": INHERITED_KEYS,
                    "OVERRIDE_KEYS": OVERRIDE_KEYS,
                    }, tree=cfg),
                "surfacer": surfacers
            })))

    except KeyError as e:
        print(e)



        if args.out:
            if os.path.isfile(args.out):
                with open(args.out, 'r') as f:
                    current = f.read().splitlines()

            content = yaml.safe_dump({
                "monitoring": {
                    "ceunix": render(ctx={
                        "INHERITED_KEYS": INHERITED_KEYS,
                        "OVERRIDE_KEYS": OVERRIDE_KEYS,
                        }, tree=cfg),
                    }
                })
            diff = list()
            for line in unified_diff(current, content.splitlines()):
                diff.append(line)
            if len(diff) > 0:
                saltret = {
                        'changed': True,
                        'comment': '\n'.join(diff)
                        }
                print(json.dumps(saltret))
            with open(args.out, 'w') as f:
                f.write(content)
                f.close()

