# -*- coding: utf-8 -*-
from setuphelpers import *

""" 
Procedures: 
https://cyberwatch.fr/actualite/cve-2021-34527-comment-identifier-et-neutraliser-la-vulnerabilite-printnightmare/

Sources: 
https://msrc.microsoft.com/update-guide/vulnerability/CVE-2021-34527
https://msrc.microsoft.com/update-guide/fr-fr/vulnerability/CVE-2021-34527

"""

# Usable WAPT package functions: install(), uninstall(), session_setup(), audit(), update_package()
# Declaring global variables - Warnings: 1) WAPT context is only available in package functions; 2) Global variables are not persistent between calls
hard_fix = False  # Will Disable the Spooler Service
kbs_dict = [
    {
        "win_vers": WindowsVersions.WindowsServer2008R2,
        "arch": "x64",
        "kbs_dl_url_list": [
            "http://download.windowsupdate.com/d/msdownload/update/software/secu/2021/07/windows6.1-kb5004953-x64_62d21485a29cad041230e4c647baeaeacc09ac7c.msu",
            "http://download.windowsupdate.com/c/msdownload/update/software/secu/2021/07/windows6.1-kb5004951-x64_2fcf9eaa66615884884cc1cb9f75fc96294cbf2a.msu",
        ],
    },
    {
        "win_vers": WindowsVersions.WindowsServer2008,
        "arch": "x64",
        "kbs_dl_url_list": [
            "http://download.windowsupdate.com/c/msdownload/update/software/secu/2021/07/windows6.0-kb5004959-x64_7bfadd426a5764d3a2886afbb73f727fae5e0f67.msu",
            "http://download.windowsupdate.com/c/msdownload/update/software/secu/2021/07/windows6.0-kb5004955-x64_b92514eb3350cddd3ce1a2e14c0cf921e3e450d2.msu",
        ],
    },
]


def install():
    # Declaring local variables
    if iswin64():
        arch = "x64"
    else:
        arch = "x86"
    patched = False

    # Installing the package
    if hard_fix:
        run_powershell("Stop-Service -Name Spooler -Force")
        run_powershell("Set-Service -Name Spooler -StartupType Disabled")

    registry_set(HKEY_LOCAL_MACHINE, r"SOFTWARE\Policies\Microsoft\Windows NT\Printers\PointAndPrint", "NoWarningNoElevationOnInstall", 0)
    registry_set(HKEY_LOCAL_MACHINE, r"SOFTWARE\Policies\Microsoft\Windows NT\Printers\PointAndPrint", "UpdatePromptSettings", 0)

    # Installing Windows KBs
    for kbs in kbs_dict:
        for kb_url in kbs["kbs_dl_url_list"]:
            kb_file = kb_url.split("/")[-1]
            kb_name = kb_file.split("-")[1].upper()
            kb_nb = kb_name.replace("KB", "")
            if kbs["arch"] == arch:
                if kbs["win_vers"] == windows_version():
                    with EnsureWUAUServRunning():
                        if not is_kb_installed(kb_nb):
                            print("Installing: {}".format(kb_file))
                            run(
                                'wusa.exe "{}" /quiet /norestart'.format(kb_file),
                                accept_returncodes=[0, 3010, 2359302, -2145124329, 2149842967],
                                timeout=3600,
                            )
                            if not is_kb_installed(kb_nb):
                                print("WARNING: {} installation do not complete".format(kb_name))
                        else:
                            print("{} is already installed".format(kb_name))
                        patched = True
    if not patched:
        print("WARNING: This PC has not been patched with a KB")


def audit():
    return_error = False
    if (
        registry_readstring(HKEY_LOCAL_MACHINE, r"SOFTWARE\Policies\Microsoft\Windows NT\Printers\PointAndPrint", "NoWarningNoElevationOnInstall")
        != "0"
        or registry_readstring(HKEY_LOCAL_MACHINE, r"SOFTWARE\Policies\Microsoft\Windows NT\Printers\PointAndPrint", "UpdatePromptSettings") != "0"
    ):
        print("ERROR: Registry values has changed")
        WAPT.write_audit_data("fix-printnightmare", "reg-values-has-changed", "yes", keep_days=365)
        print("INFO: Switching registry values to desired ones")
        registry_set(HKEY_LOCAL_MACHINE, r"SOFTWARE\Policies\Microsoft\Windows NT\Printers\PointAndPrint", "NoWarningNoElevationOnInstall", 0)
        registry_set(HKEY_LOCAL_MACHINE, r"SOFTWARE\Policies\Microsoft\Windows NT\Printers\PointAndPrint", "UpdatePromptSettings", 0)
        return_error = True

    if hard_fix:
        service_start_mode = get_service_start_mode("Spooler")
        if service_start_mode != "Disabled":
            print("ERROR: Spooler service was not disabled")
            run_powershell("Set-Service -Name Spooler -StartupType Disabled")
            return_error = True
        WAPT.write_audit_data("fix-printnightmare", "spooler-start-mode", service_start_mode, keep_days=365)
        if not run_powershell("Get-Service -Name Spooler | Select-Object -ExpandProperty Status", output_format="text").startswith("Stopped"):
            print("ERROR: Spooler service was running")
            run_powershell("Stop-Service -Name Spooler -Force")
            return_error = True
            WAPT.write_audit_data("fix-printnightmare", "spooler-was-running", "yes", keep_days=365)

    if return_error:
        return "ERROR"
    else:
        return "OK"


def update_package():
    # Declaring local variables
    result = False
    proxies = get_proxies()
    if not proxies:
        proxies = get_proxies_from_wapt_console()
    version = control.get_software_version()

    # Downloading Windows KBs
    for kbs in kbs_dict:
        for kb_url in kbs["kbs_dl_url_list"]:
            kb_file = kb_url.split("/")[-1]
            kb_name = kb_file.split("-")[1].upper()
            kb_nb = kb_name.replace("KB", "")
            if not isfile(kb_file):
                print("Downloading: %s" % kb_file)
                wget(kb_url, kb_file, proxies=proxies)

    # Changing version of the package
    if Version(version) > Version(control.get_software_version()):
        print("Software version updated (from: %s to: %s)" % (control.get_software_version(), Version(version)))
        result = True
    control.version = "%s-%s" % (Version(version), control.version.split("-", 1)[-1])
    # control.set_software_version(Version(version))
    control.save_control_to_wapt()

    # Validating update-package-sources
    return result


def is_kb_installed(hotfixid):
    r"""Check whether the Windows KB is installed

    Returns:
        boolean

    >>> is_kb_installed('3216755')
    True
    >>> is_kb_installed('Kb3216755')
    True

    """
    if not hotfixid.upper().startswith("KB"):
        hotfixid = "KB" + hotfixid
    installed_update = installed_windows_updates()
    if [kb for kb in installed_update if kb["HotFixID"].upper() == hotfixid.upper()]:
        return True
    return False