import sys
import os
import subprocess
import datetime

# ============================================================
# CONFIGURATION
# ============================================================
time_to_enable = 3600  # Inactivity timeout in seconds (3600 = 60 mins)
LOCK_SCREEN = 0        # 0 = Enable lock screen | 1 = Disable lock screen


FALLBACK_SCREENSAVER = "scrnsave.scr"


SYSTEM_SID_PREFIXES = (
    "S-1-5-18",   # SYSTEM
    "S-1-5-19",   # LOCAL SERVICE
    "S-1-5-20",   # NETWORK SERVICE
)

# ============================================================
# LOGGING
# ============================================================
LOG_FILE = "C:\\Windows\\Temp\\setup_lockscreen.log"

def log(level, msg):
    ts = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    line = ts + " [" + level + "] " + str(msg)
    print line
    try:
        f = open(LOG_FILE, "a")
        f.write(line + "\n")
        f.close()
    except Exception:
        pass


PS_EXE = r"C:\Windows\SysNative\WindowsPowerShell\v1.0\powershell.exe"
if not os.path.exists(PS_EXE):
    PS_EXE = r"C:\Windows\System32\WindowsPowerShell\v1.0\powershell.exe"


def run_ps(command):
    """Run a PowerShell command. Returns (returncode, stdout, stderr)."""
    proc = subprocess.Popen(
        [PS_EXE, "-NonInteractive", "-NoProfile", "-ExecutionPolicy", "Bypass", "-Command", command],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE
    )
    stdout, stderr = proc.communicate()
    return proc.returncode, stdout.strip(), stderr.strip()


def escape_ps(value):
    """Escape a value for embedding inside a PowerShell single-quoted string."""
    return str(value).replace("'", "''")



def set_hklm_dword(reg_path, name, value):
    ps = (
        "$k = [Microsoft.Win32.Registry]::LocalMachine.CreateSubKey('{path}'); "
        "$k.SetValue('{name}', [int]{value}, "
        "[Microsoft.Win32.RegistryValueKind]::DWord); "
        "$k.Close()"
    ).format(path=escape_ps(reg_path), name=escape_ps(name), value=int(value))

    rc, stdout, stderr = run_ps(ps)
    if rc != 0:
        log("ERROR", "HKLM\\" + reg_path + "\\" + name + " FAILED => " + stderr)
        return False
    log("OK", "HKLM\\" + reg_path + "\\" + name + " = " + str(value))
    return True


def get_hku_string(sid, reg_path, name):
    ps = (
        "New-PSDrive -PSProvider Registry -Name HKU -Root HKEY_USERS "
        "  -ErrorAction SilentlyContinue | Out-Null; "
        "try {{ "
        "  $k = [Microsoft.Win32.Registry]::Users.OpenSubKey('{sid}\\{path}'); "
        "  if ($k) {{ $k.GetValue('{name}', '') }} else {{ '' }} "
        "}} catch {{ '' }}"
    ).format(sid=escape_ps(sid), path=escape_ps(reg_path), name=escape_ps(name))

    rc, stdout, stderr = run_ps(ps)
    if rc == 0:
        return stdout.strip()
    return ""


def set_hku_string(sid, reg_path, name, value):
    ps = (
        "New-PSDrive -PSProvider Registry -Name HKU -Root HKEY_USERS "
        "  -ErrorAction SilentlyContinue | Out-Null; "
        "$k = [Microsoft.Win32.Registry]::Users.CreateSubKey('{sid}\\{path}'); "
        "$k.SetValue('{name}', '{value}', "
        "[Microsoft.Win32.RegistryValueKind]::String); "
        "$k.Close()"
    ).format(
        sid=escape_ps(sid),
        path=escape_ps(reg_path),
        name=escape_ps(name),
        value=escape_ps(value)
    )

    rc, stdout, stderr = run_ps(ps)
    if rc != 0:
        log("ERROR", "HKU\\" + sid + "\\" + reg_path + "\\" + name + " FAILED => " + stderr)
        return False
    log("OK", "HKU\\" + sid + "\\" + reg_path + "\\" + name + " = " + str(value))
    return True


def is_hive_loaded(sid):
    ps = (
        "New-PSDrive -PSProvider Registry -Name HKU -Root HKEY_USERS "
        "  -ErrorAction SilentlyContinue | Out-Null; "
        "if ([Microsoft.Win32.Registry]::Users.OpenSubKey('{sid}')) {{ 'yes' }} else {{ 'no' }}"
    ).format(sid=escape_ps(sid))
    rc, stdout, stderr = run_ps(ps)
    return stdout.strip().lower() == "yes"


def load_hive(sid, ntuser_path):
    rc, stdout, stderr = run_ps(
        "reg load 'HKU\\{sid}' '{path}'".format(
            sid=escape_ps(sid), path=escape_ps(ntuser_path)
        )
    )
    if rc != 0:
        log("WARN", "Could not load hive for SID " + sid + " from " + ntuser_path + " => " + stderr)
        return False
    log("INFO", "Loaded hive for SID " + sid + " from " + ntuser_path)
    return True


def unload_hive(sid):

    run_ps("[GC]::Collect(); [GC]::WaitForPendingFinalizers()")
    rc, stdout, stderr = run_ps(
        "reg unload 'HKU\\{sid}'".format(sid=escape_ps(sid))
    )
    if rc != 0:
        log("WARN", "Could not unload hive for SID " + sid + " => " + stderr)
        return False
    log("INFO", "Unloaded hive for SID " + sid)
    return True

# ============================================================
# PROFILE ENUMERATION
# ============================================================

def get_all_user_profiles():
    ps = (
        "$base = 'HKLM:\\SOFTWARE\\Microsoft\\Windows NT\\CurrentVersion\\ProfileList'; "
        "Get-ChildItem $base | ForEach-Object { "
        "  $sid = $_.PSChildName; "
        "  $path = (Get-ItemProperty $_.PSPath).ProfileImagePath; "
        "  Write-Output ($sid + '|' + $path) "
        "}"
    )
    rc, stdout, stderr = run_ps(ps)
    if rc != 0 or not stdout:
        log("WARN", "Could not enumerate profiles: " + stderr)
        return []

    profiles = []
    for line in stdout.splitlines():
        line = line.strip()
        if not line or "|" not in line:
            continue
        sid, profile_path = line.split("|", 1)
        sid = sid.strip()
        profile_path = profile_path.strip()

        skip = False
        for prefix in SYSTEM_SID_PREFIXES:
            if sid.startswith(prefix):
                skip = True
                break
        if skip:
            log("INFO", "Skipping system SID: " + sid)
            continue

        if not os.path.isdir(profile_path):
            log("WARN", "Profile folder missing, skipping SID " + sid + " (" + profile_path + ")")
            continue

        profiles.append((sid, profile_path))

    return profiles


def screensaver_needs_setting(current_value):

    if not current_value:
        return True

    stripped = current_value.strip().lower()
    if stripped in ("", "(none)", "none"):
        return True


    if os.sep not in current_value and "/" not in current_value:
        system32 = os.path.join(
            os.environ.get("SystemRoot", "C:\\Windows"), "System32"
        )
        resolved = os.path.join(system32, current_value)
        if os.path.exists(resolved):
            return False

        log("WARN", "Screensaver '" + current_value + "' not found in System32 - will replace")
        return True


    if not os.path.exists(current_value):
        log("WARN", "Configured screensaver not found on disk: " + current_value + " - will replace")
        return True

    return False

# ============================================================
# PER-USER SCREENSAVER CONFIG
# ============================================================

def configure_user_hive(sid, desktop_path):

    failures = 0

    current_scrnsave = get_hku_string(sid, desktop_path, "SCRNSAVE.EXE")
    log("INFO", "  Current SCRNSAVE.EXE: '" + current_scrnsave + "'")

    if screensaver_needs_setting(current_scrnsave):
        log("INFO", "  Setting blank screensaver (" + FALLBACK_SCREENSAVER + ")")
        if not set_hku_string(sid, desktop_path, "SCRNSAVE.EXE", FALLBACK_SCREENSAVER):
            failures += 1
    else:
        log("INFO", "  Screensaver already valid (" + current_scrnsave + ") - leaving unchanged")

    if not set_hku_string(sid, desktop_path, "ScreenSaveTimeOut", str(time_to_enable)):
        failures += 1
    if not set_hku_string(sid, desktop_path, "ScreenSaveActive", "1"):
        failures += 1
    if not set_hku_string(sid, desktop_path, "ScreenSaverIsSecure", "1"):
        failures += 1

    return failures

# ============================================================
# MAIN
# ============================================================

def main():
    # Validate configuration
    if not isinstance(time_to_enable, int) or time_to_enable <= 0:
        log("ERROR", "Invalid time_to_enable: " + str(time_to_enable) + " - must be a positive integer")
        sys.exit(2)
    if LOCK_SCREEN not in (0, 1):
        log("ERROR", "Invalid LOCK_SCREEN: " + str(LOCK_SCREEN) + " - must be 0 or 1")
        sys.exit(2)

    log("INFO", "=== Lock Screen / Screensaver Setup Starting ===")
    log("INFO", "LOCK_SCREEN        = " + str(LOCK_SCREEN) + "  (0=enabled, 1=disabled)")
    log("INFO", "Inactivity timeout = " + str(time_to_enable) + " seconds (" + str(time_to_enable / 60) + " minutes)")

    total_failures = 0
    desktop_path = r"Control Panel\Desktop"

    # ----------------------------------------------------------
    # Step 1: HKLM NoLockScreen policy
    # ----------------------------------------------------------
    log("INFO", "--- Step 1: HKLM NoLockScreen policy ---")
    log("WARN", "NoLockScreen only enforced on Enterprise/Education - Pro builds may ignore it")
    if not set_hklm_dword(
        r"SOFTWARE\Policies\Microsoft\Windows\Personalization",
        "NoLockScreen",
        LOCK_SCREEN
    ):
        total_failures += 1

    # ----------------------------------------------------------
    # Step 2: Enumerate all real user profiles
    # ----------------------------------------------------------
    log("INFO", "--- Step 2: Enumerating local user profiles ---")
    profiles = get_all_user_profiles()

    if not profiles:
        log("WARN", "No user profiles found - skipping per-user configuration")
    else:
        log("INFO", "Found " + str(len(profiles)) + " user profile(s) to configure")

    # ----------------------------------------------------------
    # Step 3: Configure each user profile
    # ----------------------------------------------------------
    for sid, profile_path in profiles:
        log("INFO", "--- Configuring SID: " + sid + " (" + profile_path + ") ---")

        hive_was_loaded_by_us = False
        already_loaded = is_hive_loaded(sid)

        if not already_loaded:
            ntuser_path = os.path.join(profile_path, "NTUSER.DAT")
            if not os.path.exists(ntuser_path):
                log("WARN", "  NTUSER.DAT not found at " + ntuser_path + " - skipping")
                continue
            if not load_hive(sid, ntuser_path):
                log("WARN", "  Could not load hive - skipping this profile")
                continue
            hive_was_loaded_by_us = True
        else:
            log("INFO", "  Hive already loaded (user is logged in)")

        failures = configure_user_hive(sid, desktop_path)
        total_failures += failures

        if hive_was_loaded_by_us:
            unload_hive(sid)

    # ----------------------------------------------------------
    # Step 4: Apply to .DEFAULT (covers future new user profiles)
    # ----------------------------------------------------------
    log("INFO", "--- Step 4: Configuring .DEFAULT hive (future users) ---")
    failures = configure_user_hive(".DEFAULT", desktop_path)
    total_failures += failures

    # ----------------------------------------------------------
    # Step 5: Refresh active user sessions
    # ----------------------------------------------------------
    log("INFO", "--- Step 5: Refreshing active user session parameters ---")
    rc, stdout, stderr = run_ps("RUNDLL32.EXE user32.dll,UpdatePerUserSystemParameters")
    if rc != 0:
        log("WARN", "UpdatePerUserSystemParameters returned non-zero - settings apply at next logon")
    else:
        log("OK", "User session parameters refreshed")

    # ----------------------------------------------------------
    # Summary
    # ----------------------------------------------------------
    log("INFO", "--- Summary ---")
    if total_failures == 0:
        log("INFO", "RESULT: All profiles configured successfully")
        log("INFO", "RESULT: Screensaver timeout = " + str(time_to_enable / 60) + " minutes, lock on resume = ON")
        log("INFO", "=== Completed with 0 failures ===")
        sys.exit(0)
    else:
        log("ERROR", "=== Completed with " + str(total_failures) + " failure(s) - check log above ===")
        sys.exit(2)


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        ts = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        line = ts + " [CRITICAL] Unhandled exception: " + str(e)
        print line
        try:
            f = open(LOG_FILE, "a")
            f.write(line + "\n")
            f.close()
        except Exception:
            pass
        sys.exit(3)