#!/usr/local/munki/munki-python
from Foundation import (CFPreferencesAppSynchronize,
                        CFPreferencesCopyAppValue,
                        CFPreferencesSetValue,
                        kCFPreferencesAnyHost,
                        kCFPreferencesCurrentUser)
import json
import os
import platform
import plistlib
import ssl
import subprocess
import urllib.request
import zlib

USER_AGENT = "Zentral/munkipreflight 0.17"
ZENTRAL_API_ENDPOINT = "https://%TLS_HOSTNAME%/public/munki/"  # set during the package build
ZENTRAL_API_SERVER_CERTIFICATE = "%TLS_SERVER_CERTS%"  # set during the package build
ZENTRAL_API_AUTH_TOKEN = "%TOKEN%"  # set during the enrollment in the postinstall script of the enrollment package

BUNDLE_ID = 'ManagedInstalls'
ADDITIONAL_HTTP_HEADERS_KEY = 'AdditionalHttpHeaders'
SN_HEADER = 'X-Zentral-Serial-Number'
UUID_HEADER = 'X-Zentral-UUID'


def get_serial_number_and_uuid():
    output = subprocess.check_output(["ioreg", "-a", "-c", "IOPlatformExpertDevice", "-d", "2"])
    ioreg_result = plistlib.loads(output)["IORegistryEntryChildren"][0]
    return ioreg_result["IOPlatformSerialNumber"], ioreg_result["IOPlatformUUID"]


def get_os_version():
    return subprocess.check_output(
        ["/usr/bin/sw_vers", "-productVersion"],
        encoding="utf-8"
    ).strip()


def get_arch():
    return "arm64" if platform.processor() == "arm" else "amd64"


def update_additional_http_headers(serial_number, uuid):
    # get current headers
    cfg_headers = CFPreferencesCopyAppValue(ADDITIONAL_HTTP_HEADERS_KEY,
                                            BUNDLE_ID)
    if cfg_headers:
        headers = dict(h.split(": ", 1) for h in cfg_headers)
    else:
        headers = {}
    headers[SN_HEADER] = serial_number
    headers[UUID_HEADER] = uuid
    # save the updated headers
    serialized_headers = ["{}: {}".format(k, v) for k, v in headers.items()]
    CFPreferencesSetValue(ADDITIONAL_HTTP_HEADERS_KEY,
                          serialized_headers,
                          BUNDLE_ID,
                          # to write in /var/root
                          kCFPreferencesCurrentUser,
                          kCFPreferencesAnyHost)
    CFPreferencesAppSynchronize(BUNDLE_ID)


# Zentral Munki API calls


def make_api_request(url, data=None):
    req = urllib.request.Request(url)
    req.add_header('User-Agent', USER_AGENT)
    req.add_header('Authorization', 'MunkiEnrolledMachine {}'.format(ZENTRAL_API_AUTH_TOKEN))
    if data:
        data = json.dumps(data)
        req.add_header('Content-Type', 'application/json')
        data = zlib.compress(data.encode("ascii"), 9)
        req.add_header('Content-Encoding', 'deflate')
    ctx = ssl.create_default_context(cafile=ZENTRAL_API_SERVER_CERTIFICATE or "/private/etc/ssl/cert.pem")
    response = urllib.request.urlopen(req, data=data, context=ctx)
    return json.load(response)


def fetch_job_details(machine_serial_number, os_version, arch):
    return make_api_request(
        "{}/job_details/".format(ZENTRAL_API_ENDPOINT.strip('/')),
        {'machine_serial_number': machine_serial_number,
         'os_version': os_version,
         'arch': arch}
    )


def save_job_details(job_details):
    job_details_path = "/usr/local/zentral/munki/job_details.plist"
    try:
        os.unlink(job_details_path)
    except FileNotFoundError:
        pass
    except Exception as e:
        print("Could not delete existing job details", job_details_path, str(e))
        return
    else:
        print("Deleted existing job details", job_details_path)
    try:
        with open(job_details_path, "wb") as f:
            plistlib.dump(job_details, f)
    except Exception as e:
        print("Could not write job details to", job_details_path, str(e))


def update_facts(job_details):
    facts_path = "/usr/local/zentral/munki/facts.plist"
    facts = {}
    try:
        with open(facts_path, "rb") as f:
            facts = plistlib.load(f)
    except FileNotFoundError:
        pass
    except Exception as e:
        print("Could not read", facts_path, str(e))
        return
    facts["zentral_tags"] = job_details.get("tags", [])
    facts["zentral_incidents"] = job_details.get("incidents", [])
    try:
        with open(facts_path, "wb") as f:
            plistlib.dump(facts, f)
    except Exception as e:
        print("Could not update", facts_path, str(e))


if __name__ == "__main__":
    serial_number, uuid = get_serial_number_and_uuid()
    os_version = get_os_version()
    arch = get_arch()
    update_additional_http_headers(serial_number, uuid)
    job_details = fetch_job_details(serial_number, os_version, arch)
    save_job_details(job_details)
    update_facts(job_details)
