#!/usr/bin/env python3
# Shebang needed this time to find the interpreter!

import requests
import argparse
import json
import time
import os
import re
import hashlib
import uuid

from pathlib import Path

# The new password store API is available starting with 2.5.0
try:
    from cmk.password_store.v1_unstable import Secret, parser_add_secret_option, resolve_secret_option, dereference_secret
except ModuleNotFoundError:
    # Just use the old one if importing the new one fails...
    from cmk.utils import password_store as legacy_pw_store
    
# The new storage API is available starting with 2.5.0
try:
    from cmk.server_side_programs.v1_unstable import Storage
except ModuleNotFoundError:
    # Silently ignore
    True

def get_cli_arguments():
    parser = argparse.ArgumentParser("agent_patchmon")
    parser.add_argument(
        "--baseurl",
        help="Specify the baseurl for your PatchMon installation.",
    )
    parser.add_argument(
        "--name",
        help="Chose whether 'id', 'hostname' or 'friendly_name' should be used to create hostname.",
    )
    parser.add_argument(
        "--maxexec",
        help="State a maximum time the special agent should take to retrieve patch details.",
    )
    parser.add_argument(
        "--list",
        help="State the interval (in seconds) host lists should be cached before retrieving them.",
    )
    parser.add_argument(
        "--check",
        help="State the interval (in seconds) patch info for a single hosts should stay cached before retrieving.",
    )
    parser.add_argument(
        "--reboot",
        help="If present, also query whether a reboot is required and ask for the reasons.",
        action="store_true",
    )
    parser.add_argument(
        "--grace",
        help="Get the package list for being able to create a grace period in the check plugin.",
        action="store_true",
    )
    try:
        # Try the 2.5.0 helper:
        parser_add_secret_option(
            parser,
            short="-s",
            long="--secret",
            help="Specify the id:token to log in to PatchMon.",
            required=True
        )
    except NameError:
        # Fall back to manually adding:
        parser.add_argument(
            "--secret-id",
            help="Reference the ID of the secret as stored in the password store.",
        )
        # Fall back to manually adding:
        parser.add_argument(
            "--secret",
            help="Specify the id:token to log in to PatchMon.",
        )
    return parser.parse_args()


def parse_uuid(raw: str) -> uuid.UUID | None:
    try:
        return uuid.UUID(raw)
    except ValueError:
        return None


def prepare_cache(keyid):
    # A directory is created for all cached data, the baseurl is sanitized from all slashes
    # to make sure no path traversal is possible
    h = hashlib.sha256()
    h.update(keyid.encode("utf-8"))
    sanitized_host = h.hexdigest()
    path = os.environ['OMD_ROOT'] + "/tmp/patchmon/" + sanitized_host
    tmpdir = Path(path)
    tmpdir.mkdir(parents=True, exist_ok=True)
    return path


def gently_exit(error_msg):
    agentstats = { "error": error_msg }
    print('<<<patchmon_server>>>')
    print(json.dumps(agentstats))
    exit(0)


def retrieve_hostlist(auth, url):
    try:
        response = requests.get(url, auth=basic)
    except Exception as err:
        gently_exit(str(err))
    if response.status_code > 200:
        gently_exit("Error retrieving host list. Got status code: " + str(response.status_code))
    j = response.json()
    return j
    

def get_hostlist_from_fs(auth, url, maxage, path):
    hostlist = path + "/hostlist.json"
    if os.path.exists(hostlist):
        if time.time() - os.path.getmtime(hostlist) < maxage:
            with open(hostlist, "r") as file:
                c = file.read()
                j = json.loads(c)
                return j
    j = retrieve_hostlist(auth, url)
    with open(hostlist, "w") as file:
        file.write(json.dumps(j))
    return j


def get_hostlist_from_storage(auth, url, maxage, storage):
    jstr = storage.read("hostlist", None)
    if jstr:
        return json.loads(jstr)
    j = retrieve_hostlist(auth, url)
    storage.unset(hostlist)
    storage.write("hostlist", json.dumps(j))
    return j


def get_hostlist(auth, url, maxage, storage, path):
    if storage:
        return get_hostlist_from_storage(auth, url, maxage, storage)
    else:
        return get_hostlist_from_fs(auth, url, maxage, path)


def retrieve_hoststat(auth, url):
    error = 'Error retrieving host statistics. Since this happened after retrieving '\
        'the host list, this is most likely caused by an outdated version of PatchMon. '
    try:
        response = requests.get(url, auth=basic)
    except Exception as err:
        gently_exit(error + str(err))
    if response.status_code > 200:
        gently_exit(error + "Got status code: " + str(response.status_code))
    j = response.json()
    return j


def retrieve_hostextra(auth, extraurl):
    error = 'Error retrieving host details. Since this happened after retrieving '\
        'the host list, this is most likely caused by an outdated version of PatchMon. '
    try:
        response = requests.get(extraurl, auth=basic)
    except Exception as err:
        gently_exit(error + str(err))
    if response.status_code > 200:
        gently_exit(error + "Got status code: " + str(response.status_code))
    j = response.json()
    return j


def retrieve_pkglist(auth, pkglisturl, oldpackages):
    error = 'Error retrieving package list. Since this happened after retrieving '\
        'the host list, this is most likely caused by an outdated version of PatchMon. '
    try:
        response = requests.get(pkglisturl, auth=basic, params={"updates_only": "true"})
    except Exception as err:
        gently_exit(error + str(err))
    if response.status_code > 200:
        gently_exit(error + "Got status code: " + str(response.status_code))
    j = response.json()["packages"]
    plist = compare_pkglist(j, oldpackages)
    return plist
    
    
def compare_pkglist(packages, oldpackages):
    t_now = time.time()
    plist = {}
    for p in packages:
        if p["name"] in oldpackages:
            pmeta = oldpackages[p["name"]]
            if oldpackages[p["name"]]["is_security_update"] is True:
                pmeta["is_security_update"] = True
        else:
            pmeta = { 
                "is_security_update": p["is_security_update"],
                "first_seen" : t_now,
            }
        plist[p["name"]] = pmeta
    return plist
        
    
def get_all_hostsmeta_from_fs(hostlist, maxage, path):
    stats = {
        "oldhosts": [],
        "missinghosts": [],
        "hostdata": {},
        "hostsskipped": 0,
    }
    for h in hostlist['hosts']:
        hostid = h['id']
        # Validating the UUID makes sure, no crafted id field can be used to
        # read and write in the site file system
        if parse_uuid(hostid):
            hostjson = path + "/" + hostid + ".json"
            if os.path.exists(hostjson):
                fileage = time.time() - os.path.getmtime(hostjson)
                with open(hostjson, "r") as file:
                    c = file.read()
                    j = json.loads(c)
                    j['fileage'] = fileage
                    j['maxage'] = maxage
                    stats['hostdata'][hostid] = j
                if fileage > maxage:
                    stats['oldhosts'].append(hostid)
            else:
                stats['missinghosts'].append(hostid)
    return stats


def get_all_hostsmeta_from_storage(hostlist, maxage, storage):
    stats = {
        "oldhosts": [],
        "missinghosts": [],
        "hostdata": {},
        "hostsskipped": 0,
    }
    for h in hostlist['hosts']:
        hostid = h['id']
        if parse_uuid(hostid):
            # storage.unset(hostid)
            jsdata = storage.read(hostid, None)
            if jsdata:
                j = json.loads(jsdata)
                fileage = time.time() - j['timestamp']
                j['fileage'] = fileage
                j['maxage'] = maxage
                stats['hostdata'][hostid] = j
                if fileage > maxage:
                    stats['oldhosts'].append(hostid)
            else:
                stats['missinghosts'].append(hostid)
    return stats


def get_all_hoststats(auth, urls, hostlist, maxage, maxexec, storage, path):
    url = urls["stats"]
    xurl = urls["extra"]
    if storage:
        stats = get_all_hostsmeta_from_storage(hostlist, maxage, storage)
    else:
        stats = get_all_hostsmeta_from_fs(hostlist, maxage, path)
    tstart = time.time()
    for hostid in stats['missinghosts'] + stats['oldhosts']:
        if time.time() - tstart < maxexec:
            statsurl = url.format(baseurl=args.baseurl, uuid=hostid)
            extraurl = xurl.format(baseurl=args.baseurl, uuid=hostid)
            j = retrieve_hoststat(auth, statsurl)
            # Retrieving extra data like reboot required needs another API call
            if args.reboot:
                x = retrieve_hostextra(auth, extraurl)
                j['needs_reboot'] = 1 if x['needs_reboot'] else 0
                j['reboot_reason'] = str(x['reboot_reason'])
            if args.grace and j['outdated_packages'] > 0:
                pkglisturl = urls["pkglist"].format(baseurl=args.baseurl, uuid=hostid)
                oldpackages = {}
                if hostid in stats['hostdata'] and "packages" in stats['hostdata'][hostid]:
                    oldpackages = stats['hostdata'][hostid]["packages"]
                x = retrieve_pkglist(auth, pkglisturl, oldpackages)
                j["packages"] = x
            else:
                j["packages"] = {}
            j['timestamp'] = time.time()
            if storage:
                storage.unset(hostid)
                storage.write(hostid, json.dumps(j))
            else:
                hostjson = path + "/" + hostid + ".json"
                with open(hostjson, "w") as file:
                    file.write(json.dumps(j))
            j['fileage'] = 0.0
            j['maxage'] = maxage
            stats['hostdata'][hostid] = j
        else:
            stats['hostsskipped'] += 1
    return stats


def get_access_token(args):
    userpass = None
    try:
        # The 2.5 method first: Checkmk cares for differentiating between plain password and store reference
        userpass = resolve_secret_option(args, "secret").reveal().split(":")
        return userpass
    except:
        # We have Checkmk 2.4 or 2.3 where the password store is an internal API
        # Here we have to differentiate between the two possible parameters
        if args.secret_id:
            store_reference = args.secret_id.split(":")[0]
            userpass = legacy_pw_store.lookup(legacy_pw_store.password_store_path(), store_reference).split(":")
            return userpass
        else:
            userpass = args.secret.split(":")
            return userpass


urls = {
    "api": "{baseurl}/api/v1/api/hosts",
    "host": "{baseurl}/hosts/{uuid}",
    "hosts": None,
    "stats": "{baseurl}/api/v1/api/hosts/{uuid}/stats",
    "extra": "{baseurl}/api/v1/api/hosts/{uuid}/system",
    "pkglist": "{baseurl}/api/v1/api/hosts/{uuid}/packages"
}

t_start = time.time()
args = get_cli_arguments()
urls["hosts"] = urls["api"].format(baseurl=args.baseurl)
userpass = get_access_token(args)    
basic = requests.auth.HTTPBasicAuth(userpass[0], userpass[1])

# The user identifier is randomly generated and thus should be globally unique.
# Since PatchMon host groups might be associated with this identifier, using it
# as identifier for the Storage means we do not need to pass a hostname. It also
# Prevents problems from using the same baseurl (potentially with different host
# groups).
#
# Since the same problem occurs for the hacky caching, use this key instead of
# baseurl from now on.

try:
    stor = Storage('patchmon', userpass[0])
    tmp_path = None
except NameError:
    stor = None
    tmp_path = prepare_cache(userpass[0])

# First retrieve the host list
hostlist = get_hostlist(basic, urls["hosts"], float(args.list), stor, tmp_path)

# Now retrieve info for all hosts
allstats = get_all_hoststats(basic, urls, hostlist, float(args.check), float(args.maxexec), stor, tmp_path)

for h in hostlist['hosts']:
    print('<<<<' + h[args.name] + '>>>>')
    print('<<<patchmon_patches>>>')
    try:
        allstats['hostdata'][h['id']]['url'] = urls["host"].format(baseurl=args.baseurl.rstrip('/'), uuid=h['id'])
        print(json.dumps(allstats['hostdata'][h['id']]))
        print('<<<labels:sep(0)>>>')
        print('{"patchmon/monitored":"true"}')
    except KeyError:
        print('{ "error": "No data" }')
    print('<<<<>>>>')

t_end = time.time()
agentstats = {
    "hoststotal": len(hostlist['hosts']),
    "hostsskipped": allstats["hostsskipped"],
    "duration": t_end - t_start
}

print('<<<patchmon_server>>>')
print(json.dumps(agentstats))
