api: arrange rundll32 a console logger

Signed-off-by: Simon Rozman <simon@rozman.si>

1 files changed, 57 insertions, 6 deletions

diff --git a/api/rundll32.c b/api/rundll32.c
index 89fbe88..46b93ec 100644
--- a/api/rundll32.c
+++ b/api/rundll32.c
@@ -7,6 +7,33 @@
 
 #if defined(_M_AMD64) || defined(_M_ARM64)
 
+// TODO: Log to Windows Event Log in production.
+
+#    ifdef _DEBUG
+
+static VOID CALLBACK
+ConsoleLogger(_In_ WINTUN_LOGGER_LEVEL Level, _In_ const WCHAR *LogLine)
+{
+    const WCHAR *Template;
+    switch (Level)
+    {
+    case WINTUN_LOG_INFO:
+        Template = L"[+] %s\n";
+        break;
+    case WINTUN_LOG_WARN:
+        Template = L"[-] %s\n";
+        break;
+    case WINTUN_LOG_ERR:
+        Template = L"[!] %s\n";
+        break;
+    default:
+        return;
+    }
+    fwprintf(stdout, Template, LogLine);
+}
+
+#    endif
+
 static BOOL ElevateToSystem(VOID)
 {
     HANDLE CurrentProcessToken, ThreadToken, ProcessSnapshot, WinlogonProcess, WinlogonToken, DuplicatedToken;
@@ -101,12 +128,37 @@ cleanup:
     return FALSE;
 }
 
+static void
+Init(_In_ BOOL ShowConsole)
+{
+#    ifdef _DEBUG
+    if (ShowConsole)
+    {
+        AllocConsole();
+        FILE *Stream;
+        _wfreopen_s(&Stream, L"CONOUT$", L"w", stdout);
+    }
+    WintunSetLogger(ConsoleLogger);
+#    else
+    UNREFERENCED_PARAMETER(ShowConsole);
+#    endif
+    ElevateToSystem();
+}
+
+static void Done(VOID)
+{
+    RevertToSelf();
+#    ifdef _DEBUG
+    _putws(L"\nPress any key to close . . .");
+    (VOID) _getwch();
+#    endif
+}
+
 __declspec(dllexport) VOID __stdcall CreateAdapter(HWND hwnd, HINSTANCE hinst, LPSTR lpszCmdLine, int nCmdShow)
 {
     UNREFERENCED_PARAMETER(hwnd);
     UNREFERENCED_PARAMETER(hinst);
     UNREFERENCED_PARAMETER(lpszCmdLine);
-    UNREFERENCED_PARAMETER(nCmdShow);
 
     int Argc;
     LPWSTR *Argv = CommandLineToArgvW(GetCommandLineW(), &Argc);
@@ -122,9 +174,9 @@ __declspec(dllexport) VOID __stdcall CreateAdapter(HWND hwnd, HINSTANCE hinst, L
         goto cleanupArgv;
     WINTUN_ADAPTER *Adapter;
     BOOL RebootRequired = FALSE;
-    ElevateToSystem();
+    Init(!!nCmdShow);
     DWORD Result = WintunCreateAdapter(Argv[2], Argv[3], Argc > 4 ? &RequestedGUID : NULL, &Adapter, &RebootRequired);
-    RevertToSelf();
+    Done();
     if (Result != ERROR_SUCCESS)
         goto cleanupArgv;
 
@@ -138,7 +190,6 @@ __declspec(dllexport) VOID __stdcall DeleteAdapter(HWND hwnd, HINSTANCE hinst, L
     UNREFERENCED_PARAMETER(hwnd);
     UNREFERENCED_PARAMETER(hinst);
     UNREFERENCED_PARAMETER(lpszCmdLine);
-    UNREFERENCED_PARAMETER(nCmdShow);
 
     int Argc;
     LPWSTR *Argv = CommandLineToArgvW(GetCommandLineW(), &Argc);
@@ -149,9 +200,9 @@ __declspec(dllexport) VOID __stdcall DeleteAdapter(HWND hwnd, HINSTANCE hinst, L
     if (FAILED(CLSIDFromString(Argv[2], &Adapter.CfgInstanceID)))
         goto cleanupArgv;
     BOOL RebootRequired = FALSE;
-    ElevateToSystem();
+    Init(!!nCmdShow);
     WintunDeleteAdapter(&Adapter, &RebootRequired);
-    RevertToSelf();
+    Done();
 
 cleanupArgv:
     LocalFree(Argv);