#!/usr/local/munki/munki-python
import json
import os
import plistlib
import shutil
import ssl
import subprocess
import urllib.request


USER_AGENT = "Zentral/munkipostinstall 0.8"
ENROLLMENT_URL = "%ENROLLMENT_URL%"
ENROLLMENT_SECRET = "%ENROLLMENT_SECRET%"
TLS_SERVER_CERTS = "%TLS_SERVER_CERTS%"
TLS_HOSTNAME = "%TLS_HOSTNAME%"
HAS_DISTRIBUTOR = "%HAS_DISTRIBUTOR%"


MUNKI_DIR = "/usr/local/munki"
ZENTRAL_MUNKI_DIR = "/usr/local/zentral/munki/"


def has_distributor():
    return HAS_DISTRIBUTOR == "YES"


def is_clean_install():
    return has_distributor()


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 enroll():
    req = urllib.request.Request(ENROLLMENT_URL)
    req.add_header("User-Agent", USER_AGENT)
    req.add_header("Content-Type", "application/json")
    serial_number, uuid = get_serial_number_and_uuid()
    data = json.dumps({"secret": ENROLLMENT_SECRET,
                       "serial_number": serial_number,
                       "uuid": uuid})
    ctx = ssl.create_default_context(cafile=TLS_SERVER_CERTS or "/private/etc/ssl/cert.pem")
    resp = urllib.request.urlopen(req, data=data.encode("ascii"), context=ctx)
    return json.load(resp)["token"]


def set_scripts_token(token):
    for phase in ("preflight", "postflight"):
        script_path = os.path.join(ZENTRAL_MUNKI_DIR, f"zentral_{phase}")
        with open(script_path, "r") as f:
            script = f.read()
        script = script.replace("%TOKEN%", token)
        with open(script_path, "w") as f:
            f.write(script)


def prepare_munki_preflight_postflight():
    # prepare preflight.d and postflight.d dir
    # save existing scripts
    # install new scripts
    for phase in ("preflight", "postflight"):
        script_path = os.path.join(MUNKI_DIR, phase)
        dir_path = os.path.join(MUNKI_DIR, "{}.d".format(phase))

        # remove old stuff if necessary
        if is_clean_install():
            if os.path.isfile(script_path):
                os.unlink(script_path)
            shutil.rmtree(dir_path, ignore_errors=True)

        # .d dir
        if not os.path.isdir(dir_path):
            os.makedirs(dir_path)

        # dir
        if os.path.isfile(script_path):
            if not os.path.islink(script_path):
                # backup previous script
                script_bckp_path = "{}.bckp".format(script_path)
                os.rename(script_path, script_bckp_path)
            else:
                os.unlink(script_path)

        # install zentral script runner
        ztl_main_script_path = os.path.join(ZENTRAL_MUNKI_DIR, phase)
        os.symlink(ztl_main_script_path, script_path)

        # install zentral script
        ztl_script_install_path = os.path.join(dir_path, "zentral")
        if os.path.exists(ztl_script_install_path):
            os.unlink(ztl_script_install_path)
        ztl_script_path = os.path.join(ZENTRAL_MUNKI_DIR, "zentral_{}".format(phase))
        os.symlink(ztl_script_path, ztl_script_install_path)

        if not has_distributor():
            # run script
            if subprocess.call([ztl_script_path, "auto"]):
                print("Could not run the Zentral", phase, "script")


def install_facts():
    facts_dir = "/usr/local/munki/conditions/facts"
    if os.path.isdir(facts_dir):
        script_path = os.path.join(facts_dir, "zentral.py")
        if is_clean_install() and os.path.isfile(script_path):
            os.unlink(script_path)
        if not os.path.isfile(script_path):
            os.symlink(os.path.join(ZENTRAL_MUNKI_DIR, "zentral_facts.py"), script_path)
        return True
    return False


def install_standalone_conditional_script():
    conditional_script_dir = "/usr/local/munki/conditions"
    if not os.path.exists(conditional_script_dir):
        os.makedirs(conditional_script_dir)
    script_path = os.path.join(conditional_script_dir, "zentral.py")
    if is_clean_install() and os.path.isfile(script_path):
        os.unlink(script_path)
    if not os.path.isfile(script_path):
        os.symlink(os.path.join(ZENTRAL_MUNKI_DIR, "zentral_conditional_items.py"), script_path)


# write zentral base url
def write_zentral_base_url():
    if subprocess.call(["defaults", "write",
                        "/Library/Preferences/io.zentral.plist",
                        "base_url", "https://{}".format(TLS_HOSTNAME)]):
        print("Could not write the base url in the Zentral plist")


if __name__ == "__main__":
    token = enroll()
    set_scripts_token(token)
    prepare_munki_preflight_postflight()
    if not install_facts():
        install_standalone_conditional_script()
    write_zentral_base_url()
