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

# 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
schtasks_to_disable = [
    r"\Microsoft\Windows Live\SOXE\Extractor Definitions Update Task",
    r"\Microsoft\Windows\Application Experience\Microsoft Compatibility Appraiser",
    r"\Microsoft\Windows\Application Experience\ProgramDataUpdater",
    r"\Microsoft\Windows\Application Experience\AitAgent",
    r"\Microsoft\Windows\Application Experience\PcaPatchDbTask",
    r"\Microsoft\Windows\Customer Experience Improvement Program\Consolidator",
    r"\Microsoft\Windows\Customer Experience Improvement Program\KernelCeipTask",
    r"\Microsoft\Windows\Customer Experience Improvement Program\UsbCeip",
    r"\Microsoft\Windows\UpdateOrchestrator\UpdateAssistant",
    r"\Microsoft\Windows\UpdateOrchestrator\UpdateAssistantCalendarRun",
]

services_to_disable = ["WerSvc", "dmwappushservice", "DiagTrack"]


def install():
    # Disabling Telemetry Scheduled tasks
    for task in schtasks_to_disable:
        if task_exists(task):
            try:
                disable_task(task)
            except:
                print("Unable to disable the task: %s" % task)

    # Disabling Telemetry Services
    for service in services_to_disable:
        if service_installed(service):
            try:
                service_disable(service)
            except:
                print("Unable to disable the service: %s" % service)

    # Uninstalling Telemetry Windows KB's
    for kb in (
        "3112343",
        "3083711",
        "3083325",
        "3080149",
        "3075853",
        "3075249",
        "3072318",
        "3068708",
        "3065988",
        "3064683",
        "3058168",
        "3050267",
        "3044374",
        "3035583",
        "3022345",
        "2976978",
        "3021917",
        "2990214",
        "2952664",
        "3081954",
        "3150513",
        "3139923",
        "3173040",
    ):
        run_notfatal("wusa /quiet /norestart /uninstall /kb:%s" % kb, accept_returncodes=[0, 2359303, 87, 112, 3010])

    # Disabling Telemetry with registry
    registry_setstring(HKEY_LOCAL_MACHINE, r"SOFTWARE\Policies\Microsoft\Windows\DataCollection", "AllowTelemetry", 0, REG_DWORD)

    # Disabling Telemetry with files
    autologger_file_path = makepath(r"C:\ProgramData\Microsoft\Diagnosis\ETLLogs\AutoLogger\AutoLogger-Diagtrack-Listener.etl")
    if isfile(autologger_file_path):
        pass
        # open(autologger_file_path,'wb').write('') # Permission denied by default

    # Uninstalling Microsoft Update Health Tools
    for uninstall in installed_softwares(name=r"\bMicrosoft Update Health Tools\b"):
        print("Removing: %s (%s)" % (uninstall["name"], uninstall["version"]))
        run(uninstall_cmd(uninstall["key"]))

    # Removing UpdateAssistant folder
    update_assistant_dir = r"C:\Windows\UpdateAssistant"
    if isdir(update_assistant_dir):
        remove_tree(update_assistant_dir)

    if Version(windows_version()) > Version("10.0"):
        # Windows 10 only part
        pass
    else:
        # Windows 7 only part
        registry_setstring(
            HKEY_LOCAL_MACHINE, r"SYSTEM\CurrentControlSet\Control\WMI\AutoLogger\AutoLogger-Diagtrack-Listener", "Start", 0, REG_DWORD
        )
        registry_setstring(HKEY_LOCAL_MACHINE, r"SYSTEM\ControlSet001\Control\WMI\AutoLogger\AutoLogger-Diagtrack-Listener", "Start", 0, REG_DWORD)
        registry_setstring(HKEY_LOCAL_MACHINE, r"SYSTEM\ControlSet002\Control\WMI\AutoLogger\AutoLogger-Diagtrack-Listener", "Start", 0, REG_DWORD)

        registry_setstring(
            HKEY_LOCAL_MACHINE,
            r"SYSTEM\CurrentControlSet\Control\WMI\Autologger\AutoLogger-Diagtrack-Listener\{DD17FA14-CDA6-7191-9B61-37A28F7A10DA}",
            "Enabled",
            0,
            REG_DWORD,
        )
        registry_setstring(
            HKEY_LOCAL_MACHINE,
            r"SYSTEM\ControlSet001\Control\WMI\Autologger\AutoLogger-Diagtrack-Listener\{DD17FA14-CDA6-7191-9B61-37A28F7A10DA}",
            "Enabled",
            0,
            REG_DWORD,
        )
        registry_setstring(
            HKEY_LOCAL_MACHINE,
            r"SYSTEM\ControlSet002\Control\WMI\Autologger\AutoLogger-Diagtrack-Listener\{DD17FA14-CDA6-7191-9B61-37A28F7A10DA}",
            "Enabled",
            0,
            REG_DWORD,
        )


def audit():
    # Declaring local variables
    result = "OK"

    # Auditing the package
    for svc in services_to_disable:
        if service_installed(svc):
            if service_is_running(svc):
                print("WARNING: %s is started but should be stopped" % svc)
                result = "WARNING"
            if service_get_start_mode(svc) != "Disabled":
                print("WARNING: %s start mode is NOT disabled" % svc)
                result = "WARNING"

    if str(registry_readstring(HKEY_LOCAL_MACHINE, r"SOFTWARE\Policies\Microsoft\Windows\DataCollection", "AllowTelemetry")) != "0":
        print("ERROR: AllowTelemetry is NOT disabled (often happens after Windows upgrades)")
        result = "ERROR"

    if result == "OK":
        print("OK: Windows Telemetry is well disabled on the system")
    else:
        print("Force reapplying the package since Telemetry is active...")
        WAPT.install(control.package, force=True)

    WAPT.write_audit_data_if_changed("disable-telemetry", "audit-status", result, max_count=10, keep_days=365)
    return result


def service_disable(service_name):
    """Disabling a service by its service name and stopping it"""
    import wmi

    c = wmi.WMI()
    for service in c.Win32_Service(Name=service_name):
        service.ChangeStartMode(StartMode="Disabled")
        service.StopService()


def service_change_start_mode(service_name, StartMode):
    """Changing start mode of a service
    Usable values: "Boot" "System" "Automatic" "Manual" "Disabled"
    Informations: https://msdn.microsoft.com/en-us/library/aa384896(v=vs.85).aspx
    """
    import wmi

    c = wmi.WMI()
    for service in c.Win32_Service(Name=service_name):
        service.ChangeStartMode(StartMode=StartMode)


def service_set_start_mode(service_name, StartMode):
    """Changing start mode of a service
    Usable values: "Boot" "System" "Automatic" "Manual" "Disabled"
    Informations: https://msdn.microsoft.com/en-us/library/aa384896(v=vs.85).aspx
    """
    import wmi

    c = wmi.WMI()
    for service in c.Win32_Service(Name=service_name):
        service.ChangeStartMode(StartMode=StartMode)


def service_get_start_mode(service_name):
    """Getting actual start mode of a service
    Returned values: "Boot" "System" "Auto" "Manual" "Disabled"
    """
    import wmi

    s = wmi.WMI()
    for service in s.Win32_Service(Name=service_name):
        # print(service.State, service.StartMode, service.Name, service.DisplayName)
        return service.StartMode