diff --git a/kvc/ControllerProcessOperations.cpp b/kvc/ControllerProcessOperations.cpp index f1e66da..3a8e658 100644 --- a/kvc/ControllerProcessOperations.cpp +++ b/kvc/ControllerProcessOperations.cpp @@ -671,79 +671,58 @@ bool Controller::GetProcessProtectionByName(const std::wstring& processName) noe bool Controller::ListProtectedProcesses() noexcept { if (!BeginDriverSession()) { - EndDriverSession(true); + EndDriverSession(true); return false; } auto processes = GetProcessList(); DWORD count = 0; - // Enable console virtual terminal processing for color output - HANDLE hConsole = GetStdHandle(STD_OUTPUT_HANDLE); - DWORD consoleMode = 0; - GetConsoleMode(hConsole, &consoleMode); - SetConsoleMode(hConsole, consoleMode | ENABLE_VIRTUAL_TERMINAL_PROCESSING); + // Enable ANSI colors using Utils function + if (!Utils::EnableConsoleVirtualTerminal()) { + ERROR(L"Failed to enable console colors"); + } - // ANSI color codes - auto GREEN = L"\033[92m"; - auto YELLOW = L"\033[93m"; - auto BLUE = L"\033[94m"; - auto HEADER = L"\033[97;44m"; - auto RESET = L"\033[0m"; - - std::wcout << GREEN; + std::wcout << Utils::ProcessColors::GREEN; std::wcout << L"\n -------+------------------------------+---------+-----------------+-----------------------+-----------------------+--------------------\n"; - std::wcout << HEADER; + std::wcout << Utils::ProcessColors::HEADER; std::wcout << L" PID | Process Name | Level | Signer | EXE sig. level | DLL sig. level | Kernel addr. "; - std::wcout << RESET << L"\n"; - std::wcout << GREEN; + std::wcout << Utils::ProcessColors::RESET << L"\n"; + std::wcout << Utils::ProcessColors::GREEN; std::wcout << L" -------+------------------------------+---------+-----------------+-----------------------+-----------------------+--------------------\n"; for (const auto& entry : processes) { if (entry.ProtectionLevel > 0) { - const wchar_t* processColor = GREEN; - - // Color coding based on signature levels - bool hasUncheckedSignatures = (entry.SignatureLevel == 0x00 || entry.SectionSignatureLevel == 0x00); + // Use centralized color logic instead of duplicated code + const wchar_t* processColor = Utils::GetProcessDisplayColor( + entry.SignerType, + entry.SignatureLevel, + entry.SectionSignatureLevel + ); - if (hasUncheckedSignatures) { - processColor = BLUE; // Blue for processes with unchecked signatures - } else { - // Check if it's a user process (non-system signer) - bool isUserProcess = (entry.SignerType != static_cast(PS_PROTECTED_SIGNER::Windows) && - entry.SignerType != static_cast(PS_PROTECTED_SIGNER::WinTcb) && - entry.SignerType != static_cast(PS_PROTECTED_SIGNER::WinSystem) && - entry.SignerType != static_cast(PS_PROTECTED_SIGNER::Lsa)); - processColor = isUserProcess ? YELLOW : GREEN; - } - - std::wcout << processColor; wchar_t buffer[512]; swprintf_s(buffer, L" %6d | %-28s | %-3s (%d) | %-11s (%d) | %-14s (0x%02x) | %-14s (0x%02x) | 0x%016llx\n", - entry.Pid, - entry.ProcessName.c_str(), - Utils::GetProtectionLevelAsString(entry.ProtectionLevel), - entry.ProtectionLevel, - Utils::GetSignerTypeAsString(entry.SignerType), - entry.SignerType, - Utils::GetSignatureLevelAsString(entry.SignatureLevel), - entry.SignatureLevel, - Utils::GetSignatureLevelAsString(entry.SectionSignatureLevel), - entry.SectionSignatureLevel, - entry.KernelAddress); - std::wcout << buffer; + entry.Pid, + entry.ProcessName.length() > 28 ? + (entry.ProcessName.substr(0, 25) + L"...").c_str() : entry.ProcessName.c_str(), + Utils::GetProtectionLevelAsString(entry.ProtectionLevel), entry.ProtectionLevel, + Utils::GetSignerTypeAsString(entry.SignerType), entry.SignerType, + Utils::GetSignatureLevelAsString(entry.SignatureLevel), entry.SignatureLevel, + Utils::GetSignatureLevelAsString(entry.SectionSignatureLevel), entry.SectionSignatureLevel, + entry.KernelAddress); + + std::wcout << processColor << buffer << Utils::ProcessColors::RESET; count++; } } - std::wcout << GREEN; + std::wcout << Utils::ProcessColors::GREEN; std::wcout << L" -------+------------------------------+---------+-----------------+-----------------------+-----------------------+--------------------\n"; - std::wcout << RESET << L"\n"; + std::wcout << Utils::ProcessColors::RESET; - SUCCESS(L"Enumerated %d protected processes", count); - - EndDriverSession(true); // Force cleanup - return true; + SUCCESS(L"Listed %d protected processes", count); + EndDriverSession(true); + return count > 0; } // ============================================================================ @@ -1022,12 +1001,18 @@ bool Controller::UnprotectBySigner(const std::wstring& signerName) noexcept { } else { INFO(L"Batch unprotection completed: %d/%d processes successfully unprotected", successCount, totalCount); } - - // ZMIANA: Użyj force cleanup tak jak w ListProcessesBySigner EndDriverSession(true); return successCount > 0; } +/** + * Lists all processes that have the specified signer type. + * Displays process information in a formatted table including PID, name, protection level, + * signer type, signature levels, and kernel address. + * + * @param signerName The name of the signer type to filter by (e.g., "Windows", "Antimalware", "WinTcb") + * @return true if processes were found and displayed, false if signer type is invalid or no processes match + */ bool Controller::ListProcessesBySigner(const std::wstring& signerName) noexcept { auto signerType = Utils::GetSignerTypeFromString(signerName); if (!signerType) { @@ -1035,63 +1020,76 @@ bool Controller::ListProcessesBySigner(const std::wstring& signerName) noexcept return false; } - // Pobierz dane PRZED operacjami konsoli std::vector processes; if (!BeginDriverSession()) { return false; } - processes = GetProcessList(); // Pobierz dane gdy sterownik aktywny - EndDriverSession(true); // Natychmiast zamknij sterownik + processes = GetProcessList(); // Collect data while driver is active + EndDriverSession(true); // Close driver session immediately - // Reszta operacji BEZ sterownika - HANDLE hConsole = GetStdHandle(STD_OUTPUT_HANDLE); - CONSOLE_SCREEN_BUFFER_INFO consoleInfo; - GetConsoleScreenBufferInfo(hConsole, &consoleInfo); - WORD originalColor = consoleInfo.wAttributes; + // Enable ANSI colors - same as ListProtectedProcesses + if (!Utils::EnableConsoleVirtualTerminal()) { + ERROR(L"Failed to enable console colors"); + // Continue anyway, just without colors + } bool foundAny = false; INFO(L"Processes with signer: %s", signerName.c_str()); + + // Use same table formatting as ListProtectedProcesses + std::wcout << Utils::ProcessColors::GREEN; + std::wcout << L"\n -------+------------------------------+---------+-----------------+-----------------------+-----------------------+--------------------\n"; + std::wcout << Utils::ProcessColors::HEADER; + std::wcout << L" PID | Process Name | Level | Signer | EXE sig. level | DLL sig. level | Kernel addr. "; + std::wcout << Utils::ProcessColors::RESET << L"\n"; + std::wcout << Utils::ProcessColors::GREEN; std::wcout << L" -------+------------------------------+---------+-----------------+-----------------------+-----------------------+--------------------\n"; - std::wcout << L" PID | Process Name | Level | Signer | EXE sig. level | DLL sig. level | Kernel addr.\n"; - std::wcout << L" -------+------------------------------+---------+-----------------+-----------------------+-----------------------+--------------------\n"; + std::wcout << Utils::ProcessColors::RESET; for (const auto& entry : processes) { if (entry.SignerType == signerType.value()) { foundAny = true; - if (entry.ProtectionLevel > 0) { - SetConsoleTextAttribute(hConsole, FOREGROUND_RED | FOREGROUND_GREEN | FOREGROUND_INTENSITY); - } + // Use centralized color logic + const wchar_t* processColor = Utils::GetProcessDisplayColor( + entry.SignerType, + entry.SignatureLevel, + entry.SectionSignatureLevel + ); wchar_t buffer[512]; - swprintf_s(buffer, L" %6d | %-28s | %s (%d) | %s (%d) | %s (0x%02x) | %s (0x%02x) | 0x%016llx\n", + + std::wcout << processColor; // Apply color + + // Use consistent column widths and formatting + swprintf_s(buffer, L" %6d | %-28s | %-3s (%d) | %-11s (%d) | %-14s (0x%02x) | %-14s (0x%02x) | 0x%016llx\n", entry.Pid, - entry.ProcessName.length() > 28 ? (entry.ProcessName.substr(0, 25) + L"...").c_str() : entry.ProcessName.c_str(), + entry.ProcessName.length() > 28 ? + (entry.ProcessName.substr(0, 25) + L"...").c_str() : entry.ProcessName.c_str(), Utils::GetProtectionLevelAsString(entry.ProtectionLevel), entry.ProtectionLevel, Utils::GetSignerTypeAsString(entry.SignerType), entry.SignerType, Utils::GetSignatureLevelAsString(entry.SignatureLevel), entry.SignatureLevel, Utils::GetSignatureLevelAsString(entry.SectionSignatureLevel), entry.SectionSignatureLevel, - entry.KernelAddress - ); + entry.KernelAddress); + std::wcout << buffer; - - SetConsoleTextAttribute(hConsole, originalColor); + std::wcout << Utils::ProcessColors::RESET; // Reset color after each line } } - std::wcout << L" -------+------------------------------+---------+-----------------+-----------------------+-----------------------+--------------------\n"; - if (!foundAny) { - INFO(L"No processes found with signer: %s", signerName.c_str()); + std::wcout << L"\nNo processes found with signer type: " << signerName << L"\n"; + return false; } - SetConsoleTextAttribute(hConsole, originalColor); - std::wcout << std::flush; + std::wcout << Utils::ProcessColors::GREEN; + std::wcout << L" -------+------------------------------+---------+-----------------+-----------------------+-----------------------+--------------------\n"; + std::wcout << Utils::ProcessColors::RESET; - return foundAny; + return true; } // ============================================================================ // PROCESS NAME-BASED OPERATIONS diff --git a/kvc/Utils.cpp b/kvc/Utils.cpp index 6a43095..2cb8258 100644 --- a/kvc/Utils.cpp +++ b/kvc/Utils.cpp @@ -888,4 +888,30 @@ namespace Utils return {}; } } -} \ No newline at end of file + // Color Functions Implementation + bool Utils::EnableConsoleVirtualTerminal() noexcept + { + HANDLE hConsole = GetStdHandle(STD_OUTPUT_HANDLE); + if (hConsole == INVALID_HANDLE_VALUE) return false; + + DWORD consoleMode = 0; + if (!GetConsoleMode(hConsole, &consoleMode)) return false; + + return SetConsoleMode(hConsole, consoleMode | ENABLE_VIRTUAL_TERMINAL_PROCESSING); + } + + const wchar_t* Utils::GetProcessDisplayColor(UCHAR signerType, UCHAR signatureLevel, UCHAR sectionSignatureLevel) noexcept + { + bool hasUncheckedSignatures = (signatureLevel == 0x00 || sectionSignatureLevel == 0x00); + if (hasUncheckedSignatures) { + return ProcessColors::BLUE; + } + + bool isUserProcess = (signerType != static_cast(PS_PROTECTED_SIGNER::Windows) && + signerType != static_cast(PS_PROTECTED_SIGNER::WinTcb) && + signerType != static_cast(PS_PROTECTED_SIGNER::WinSystem) && + signerType != static_cast(PS_PROTECTED_SIGNER::Lsa)); + + return isUserProcess ? ProcessColors::YELLOW : ProcessColors::GREEN; + } +} diff --git a/kvc/Utils.h b/kvc/Utils.h index cc836a1..704055f 100644 --- a/kvc/Utils.h +++ b/kvc/Utils.h @@ -87,4 +87,16 @@ namespace Utils std::vector& second) noexcept; std::vector DecryptXOR(const std::vector& encryptedData, const std::array& key) noexcept; -} \ No newline at end of file + + // Console coloring utilities for process display + struct ProcessColors { + static constexpr const wchar_t* GREEN = L"\033[92m"; + static constexpr const wchar_t* YELLOW = L"\033[93m"; + static constexpr const wchar_t* BLUE = L"\033[94m"; + static constexpr const wchar_t* HEADER = L"\033[97;44m"; + static constexpr const wchar_t* RESET = L"\033[0m"; + }; + + bool EnableConsoleVirtualTerminal() noexcept; + const wchar_t* GetProcessDisplayColor(UCHAR signerType, UCHAR signatureLevel, UCHAR sectionSignatureLevel) noexcept; +}