#!/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 sys
import yaml
import json
import argparse
from urllib.parse import urlparse

import axconfig, onepdns, re
from collections import OrderedDict

config_footer = """
surfacer {
  type: PROMETHEUS

  ignore_metrics_with_label {
    key: "probe",
    value: "sysvars",
  }
}
"""
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
    "product",
    "datacenter"
)


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(l: list, keyfield: str="key") -> dict:
    if isinstance(l, list):
        return {i[keyfield]: i["value"] for i in l }
    return l


# Adding labels into probe config
#
#   probe {
#     ...
# +   additional_label {
# +     key:   {key}
# +     value: {val}
# +   }
#   }
#
def render_labels(labels: dict):
    if not 'datacenter' 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": {
#                   "query_type": "A",
                   "resolved_domain": url.path.lstrip('/')
               }
           }


def custom_params(url, proto):
    opts = {
        "port": str(url.port),
        "proto": url.scheme,
        "target": "@target@"
    }
    x = {**proto}
#    print(x)
    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 resource(*args, **kwargs):
#    return **kwargs

# TODO: rename thos function
# search hosts by mask in powerdns, apply regex filters
# and save result in `path`.json file
def host_search_file(mask, filters, path):
    r = {}
    r["resource"] = list()
    cfg = axconfig.load_config(path='axc_driver.conf', keypath='pdns')
    #cfgext = axconfig.merge_with_subpath(cfg, 'map:ext')
    pdns = onepdns.PDNS(**cfg)
    records = pdns.search(mask)
    records = onepdns.filter_records(records, "type", ["A", "CNAME"])
    for record in records:
        if filters and 'host.match' in filters:
            m = re.match(r'^.*' + filters['host.match'] + r'.*$', record['name'])
            if m:
                tmp = {}
                fqdn = onepdns.dotless_fqdn(f"{record['name']}")
                dc = datacenter(fqdn)
                tmp["name"] = fqdn
                if dc:
                    tmp["labels"] = {}
                    tmp["labels"]["datacenter"] = dc
                r["resource"].append(tmp)

        else:
                tmp = {}
                fqdn = onepdns.dotless_fqdn(f"{record['name']}")
                dc = datacenter(fqdn)
                tmp["name"] = fqdn
                if dc:
                    tmp["labels"] = {}
                    tmp["labels"]["datacenter"] = dc
                r["resource"].append(tmp)
    with open(path, 'w') as f:
        f.write(json.dumps(r, indent=2))


def http_validators(filters):
    i=1
    if filters and "status" in filters:
        return {
            "validator": [
                {
                    "name": f"status_{i}",
                    "http_validator": {
                        "success_status_codes": filters["status"]
                    }
                }
            ]
        }
    return {
            "validator": [
                {
                    "name": f"status_{i}",
                    "http_validator": {
                        "success_status_codes": "200-299"
                    }
                }
            ]
    }
#
#  Parse rule
#
def parse_rule(rule :str, __vars :dict, name :str, targets :dict, proto: dict):  # -> ???
    # 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:
        type = "EXTERNAL"
    else:
        type = url.scheme.upper()

    # Check if the hostname is a ${variable} which points to the `targets` record
    # Only one variable is supported
    m = re.match(r'^\$\{(.+)\}$', url.hostname)
    if m:
        # Lookup ${variable} in `targets` dict,
        # which was specified in some upper level of the yaml config tree
        # and save it into `target`
        target = targets[ m.group(1) ]
# TODO:
#    - validators needed

        # merge `target`-Dict  and `filters`-ODict into single OrderedDict
        # Since there is no '+' operator for 2 OrderedDict it's passed as
        # sum of list of items into new OrderedDict constructor
        filters = OrderedDict(  list( OrderedDict(target).items() )  +  list( filters.items() )  )
        # and pass `mask`, `filter`, `outfile`
        host_search_file(target["host.mask"], filters, f'{name}.json')

    # target hostname isn't variable
    else:
        host_search_file(url.hostname, filters, f'{name}.json')
    r = {
        "type": type,
        "targets": {
            "file_targets": {
                "file_path": f"{name}.json"
            }
        }
    }

    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:
        p = custom_params(url, proto[url.scheme])

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





def render_alerts(alerts: dict, teams: dict, severity: str):
    import copy
    __providers = copy.deepcopy(alerts['providers'])
#    if "http_notify" in __providers and "data" in __providers["http_notify"]:
#        d = str(__providers["http_notify"]["data"])
##        print(d)
##        d = " ".join(d)
#        d = d.replace('\n', ' ')
#        __providers["http_notify"]["data"] = list(d) # list(__providers["http_notify"]["data"])
    __rules = alerts['rules']
    r = list()
    for rule in __rules:
        if "condition" in rule and 2 == len(rule["condition"].split('/')):
            failures, total = rule["condition"].split('/')
            p2send = list()
            if "default" in rule['notify']:
                p2send = rule['notify']["default"].copy()
            if severity in rule['notify']:
                for x in rule['notify'][severity]:
                    if x not in p2send:
                        p2send.append(x)
            for p in p2send:
                tmp = {
                    "condition": {
                        "failures": int(failures),
                        "total": int(total)
                    },
                    "notify": {
                        p: dict(__providers[p])
                    }
                }
                r.append(tmp)

    return { "alert": r }

#
# Render "probe { ... }" for the service
#
def render_svcs(svcs: list, labels: dict, targets: dict, proto: dict, config:dict, teams: dict, alerts: dict) -> str:
    r = list()
    for svc in svcs:
        __labels = {**labels}
        __config = {**config}
        __teams = {**teams}
        __alerts = {**alerts}
        __vars = {}
        name = None
        if 'name' in svc:
            name = svc['name']
        if 'config' in svc:
            __config = {**config, **svc["config"]}
        if 'alerts' in svc:
            __alerts = {**alerts, **svc["alerts"]}
        if 'labels' in svc:
            for label in svc['labels']:
                if svc['labels'][label].startswith('$') and svc['labels'][label] in __vars:
                    __labels[label] = __vars[ svc['labels'][label] ]
                else:
                    __labels[label] = svc['labels'][label]
        if 'rule' in svc:
            rule_content = parse_rule(svc['rule'], __vars, svc['name'], targets, proto)

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

        tmp = {
            "name": name,
            "interval": "60s",
            "timeout": "3s"
        }
#        tmp = {**tmp, **__config}
        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(tree: dict, labels: dict, targets: dict, proto: dict, config: dict, teams: dict, alerts: dict) -> str:
    # local _labels/_targets dict.
    # We can easily add labels/targets from current subtree
    # Those will be reverted back when we return from recursion
    local_labels = {**labels}
    local_targets = {**targets}
    local_proto = proto_fmt(proto)
    local_config = {**config}
    local_teams = {**teams}
    local_alerts = {**alerts}
    r = list()
    if   'labels' in tree:
        local_labels = {**labels, **tree['labels']}

    if 'targets' in tree:
        local_targets = {**targets, **tree['targets']}

    if 'proto' in tree:
        local_proto = {**proto, **proto_fmt(tree['proto'])}

    if 'config' in tree:
        local_config = {**config, **tree['config']}

    if 'teams' in tree:
        local_teams = {**teams, **tree['teams']}

    if 'alerts' in tree:
        local_alerts = {**alerts, **tree['alerts']}

    for entry in tree:
        if  entry in ['labels', 'targets', 'proto', 'config', 'teams', 'alerts']:
            pass

        elif 'services' == entry:
            r += render_svcs(tree['services'], local_labels, local_targets, local_proto, local_config, local_teams, local_alerts)

        else: # it's name of subtree, parse it!   pass local labels and targets
            r += render(tree[entry], local_labels, local_targets, local_proto, local_config, local_teams, local_alerts)
    return r

def print_textpb(obj, level=0, indent=2, upkey=""):
    if isinstance(obj, dict):
        for k in obj:
            if isinstance(obj[k], list):
                for v in obj[k]:
                    print(' '*indent*level + f'{k}' + ' {')
                    print_textpb(v, level=level+1, upkey=k)
                    print(' '*indent*level + '}')
            elif isinstance(obj[k], dict):
                print(' '*indent*level + f'{k}' + ' {')
                print_textpb(obj[k], level=level+1, upkey=k)
                print(' '*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')
                print(' '*indent*level + f'{k}: "{s}"')
            else:
                print(' '*indent*level + f'{k}: {obj[k]}')

def parse_args():
    parser = argparse.ArgumentParser( prog = 'cbhelper',
                                      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 = json.loads(stdinbuf)
        cfg = yaml.safe_load(stdinbuf)
        if args.out:
            with open(args.out, 'w') as f:
                f.write(render(cfg, labels={}, targets={}))
                f.write(config_footer)
        else:
            print_textpb({
                "probe": render(cfg, labels={}, targets={}, proto={}, config={}, teams={}, alerts={}),
                "surfacer": surfacers
            })
#            print(yaml.dump({
#                "probe": render(cfg, labels={}, targets={}, proto={}, config={}, teams={}, alerts={}),
#                "surfacer": surfacers
#            }, indent=2))
#            print(json.dumps({
#                "probe": render(cfg, labels={}, targets={}, proto={}, config={}, teams={}, alerts={}),
#                "surfacer": surfacers
#            }, indent=2))

    except KeyError as e:
        print(e)

