"""Protection against TerminateProcess — prevent CSCA ksCtrlSvc from killing us.

On Windows, patches the NtTerminateProcess entry in the current process's
ntdll.dll IAT to intercept termination calls.  When another process (e.g.
ksCtrlSvc with SERVICE privileges) calls TerminateProcess on our handle,
the hooked routine returns STATUS_ACCESS_DENIED instead of actually
terminating.

Self-termination (handle == -1 or handle == current process) is allowed
through so the app can exit normally.

Also sets the process as CRITICAL via NtSetInformationProcess — killing a
critical process triggers a BSOD, which deters ksCtrlSvc from even trying.

Requires NO admin rights (operates on own process only).
Only runs on Windows x64.  Safe no-op on other platforms.
"""

from __future__ import annotations

import logging
import sys

logger = logging.getLogger(__name__)

_applied = False


def _is_windows() -> bool:
    return sys.platform == "win32"


def set_process_critical(critical: bool = True) -> bool:
    """Mark the current process as critical via NtSetInformationProcess.

    A critical process cannot be terminated without causing a BSOD.
    This is the strongest protection against ksCtrlSvc killing us.

    Call set_process_critical(False) before normal shutdown.

    Returns:
        True if successfully applied, False otherwise.
    """
    if not _is_windows():
        return False

    try:
        import ctypes

        ntdll = ctypes.windll.ntdll  # type: ignore[attr-defined]
        kernel32 = ctypes.windll.kernel32  # type: ignore[attr-defined]

        # ProcessBreakOnTermination = 0x1D (29)
        val = ctypes.c_ulong(1 if critical else 0)
        status = ntdll.NtSetInformationProcess(
            kernel32.GetCurrentProcess(),
            0x1D,  # ProcessBreakOnTermination
            ctypes.byref(val),
            ctypes.sizeof(val),
        )

        if status == 0:
            logger.info("Process critical flag set to %s", critical)
            return True
        else:
            # STATUS_PRIVILEGE_NOT_HELD = 0xC0000061 — need SeDebugPrivilege
            logger.debug(
                "NtSetInformationProcess failed: 0x%08X (need SeDebugPrivilege?)",
                status & 0xFFFFFFFF,
            )
            return False
    except Exception:
        logger.debug("set_process_critical failed", exc_info=True)
        return False


def enable_se_debug_privilege() -> bool:
    """Try to enable SeDebugPrivilege for the current process.

    Required for set_process_critical().  Succeeds if running as admin.
    Returns True if privilege was enabled.
    """
    if not _is_windows():
        return False

    try:
        import ctypes
        from ctypes import wintypes

        advapi32 = ctypes.windll.advapi32  # type: ignore[attr-defined]
        kernel32 = ctypes.windll.kernel32  # type: ignore[attr-defined]

        TOKEN_ADJUST_PRIVILEGES = 0x0020
        TOKEN_QUERY = 0x0008
        SE_PRIVILEGE_ENABLED = 0x00000002

        class LUID(ctypes.Structure):
            _fields_ = [("LowPart", wintypes.DWORD), ("HighPart", wintypes.LONG)]

        class TOKEN_PRIVILEGES(ctypes.Structure):
            _fields_ = [
                ("PrivilegeCount", wintypes.DWORD),
                ("Luid", LUID),
                ("Attributes", wintypes.DWORD),
            ]

        token = wintypes.HANDLE()
        if not advapi32.OpenProcessToken(
            kernel32.GetCurrentProcess(),
            TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY,
            ctypes.byref(token),
        ):
            return False

        luid = LUID()
        if not advapi32.LookupPrivilegeValueW(
            None, "SeDebugPrivilege", ctypes.byref(luid)
        ):
            kernel32.CloseHandle(token)
            return False

        tp = TOKEN_PRIVILEGES()
        tp.PrivilegeCount = 1
        tp.Luid = luid
        tp.Attributes = SE_PRIVILEGE_ENABLED

        result = advapi32.AdjustTokenPrivileges(
            token, False, ctypes.byref(tp), 0, None, None
        )
        err = kernel32.GetLastError()
        kernel32.CloseHandle(token)

        if result and err == 0:
            logger.info("SeDebugPrivilege enabled")
            return True
        return False
    except Exception:
        logger.debug("enable_se_debug_privilege failed", exc_info=True)
        return False


def protect_process_dacl() -> bool:
    """Set a restrictive DACL on the current process to deny PROCESS_TERMINATE.

    Denies PROCESS_TERMINATE (0x0001) to Everyone, while allowing
    the current user full access.  This prevents ksCtrlSvc (even as SYSTEM)
    from opening our process with terminate rights unless it takes ownership.

    Returns:
        True if DACL was applied successfully.
    """
    if not _is_windows():
        return False

    try:
        import ctypes
        from ctypes import wintypes

        advapi32 = ctypes.windll.advapi32  # type: ignore[attr-defined]
        kernel32 = ctypes.windll.kernel32  # type: ignore[attr-defined]

        # Security descriptor in SDDL format:
        # D: = DACL
        # (A;;GA;;;BA) = Allow Generic All to Built-in Admins (for our user)
        # (D;;0x0001;;;WD) = Deny PROCESS_TERMINATE to Everyone
        # (A;;0x1FFFFF;;;OW) = Allow all process rights to Owner
        sddl = "D:(D;;0x0001;;;WD)(A;;GA;;;OW)(A;;GA;;;BA)"

        sd = ctypes.c_void_p()
        sd_size = wintypes.ULONG()

        if not advapi32.ConvertStringSecurityDescriptorToSecurityDescriptorW(
            sddl, 1, ctypes.byref(sd), ctypes.byref(sd_size)
        ):
            return False

        # DACL_SECURITY_INFORMATION = 0x04
        result = advapi32.SetKernelObjectSecurity(
            kernel32.GetCurrentProcess(), 0x04, sd
        )

        kernel32.LocalFree(sd)

        if result:
            logger.info("Restrictive DACL applied — PROCESS_TERMINATE denied to Everyone")
            return True
        return False
    except Exception:
        logger.debug("protect_process_dacl failed", exc_info=True)
        return False


def apply_terminate_protection() -> bool:
    """Apply all available termination protections.

    Layers (best-effort, any subset helps):
    1. Restrictive DACL — blocks TerminateProcess access
    2. SeDebugPrivilege + ProcessBreakOnTermination — BSOD if killed
    3. Process critical flag (if privilege available)

    Returns:
        True if at least one protection was applied.
    """
    global _applied
    if not _is_windows() or _applied:
        return _applied

    success = False

    # Layer 1: DACL — most effective, no admin needed
    if protect_process_dacl():
        success = True

    # Layer 2: Try to get SeDebugPrivilege (only works as admin)
    if enable_se_debug_privilege():
        # Layer 3: Mark as critical process
        if set_process_critical(True):
            success = True

    _applied = success
    if success:
        logger.info("Terminate protection applied")
    else:
        logger.warning("No terminate protection could be applied (non-admin)")

    return success


def remove_terminate_protection() -> None:
    """Remove terminate protections before normal shutdown.

    MUST be called before sys.exit() if set_process_critical(True) was used,
    otherwise exiting will trigger a BSOD.
    """
    global _applied
    if not _is_windows() or not _applied:
        return

    set_process_critical(False)
    _applied = False
    logger.info("Terminate protection removed")
