sku/thesku

/*
    AUTHOR:         sku/thesku [http://www.ownedcore.com/forums/members/69674-sku.html]

    CREDITS:        shadowdancer_ on irc.synirc.net#d3dev
                  #d3dev [irc.synirc.net]
                  #d3se [irc.rizon.net]

    DESCRIPTION:    Logs all login-relevant protobuf message constructor calls [x86 only]

    VERSION:        1.2

    DEPENDENCIES:   Boost [http://www.boost.org]
                  HadesMem [http://code.google.com/p/hadesmem/]
                  AsmJit [http://code.google.com/p/asmjit/]
                  BeaEngine [http://www.beaengine.org/]

    TODO:           Add all constructors
                  Detour WSARecv/WSASend

    LEGAL:          Code released to public domain by sku, no copyright
                  Dependencies are copyrighted and licensed by their respective owners
                  Use at your own risk

    CHANGELOG:      1.1 Added timestamps to log messages.
                  1.2   Removed duplicate code
*/

// Windows.
#include <Windows.h>

// Standard.
#include <algorithm>
#include <ctime>
#include <fstream>
#include <memory>
#include <stdexcept>
#include <tuple>
#include <vector>

// Boost.
#include <boost/exception/all.hpp>
#include <boost/format.hpp>

// HadesMem (v3, old).
#define ASMJIT_API
#include <HadesMemory/MemoryMgr.hpp>
#include <HadesMemory/Patcher.hpp>

// Link libraries.
#pragma comment(lib, "libAsmJit")
#pragma comment(lib, "libBeaEngine")
#pragma comment(lib, "libMemory")
#pragma comment(lib, "psapi")

// Forward declarations.
void __stdcall Log(std::string const& message);
void __stdcall Log(boost::format const& fmt);
void __cdecl AddConstructorAddresses();
void __cdecl CtorTrampoline();
DWORD __stdcall CtorHandler(DWORD messagePtr, DWORD ctorAddress);

// Constants.
DWORD const kIDA_BASE = 0x3C910000;
std::string const kLOG_FILE_PATH = "dump-msg-ctor-calls.log";

// Typedefs.
typedef std::tuple<std::shared_ptr<Hades::Memory::PatchRaw>, DWORD> PatchTuple;
typedef std::tuple<DWORD, std::string> CtorTuple;
typedef std::runtime_error Error;

// Globals.
std::ofstream g_logFile;
std::vector<PatchTuple> g_patches;
std::vector<CtorTuple> g_messageCtors;
Hades::Memory::MemoryMgr g_memory(GetCurrentProcessId());


// Dll entry point.
BOOL WINAPI DllMain(HINSTANCE, DWORD, LPVOID) {
    return TRUE;
}

// Dll injector initialize export.
extern "C" __declspec(dllexport) void Initialize() {
    try {
        try {
            // Delete log file if it exists already.
            std::remove(kLOG_FILE_PATH.c_str());
        } catch (...) {
            // Exception may be ignored.
        }

        // For debugging purposes.
        Log("Initialize called.");

        // Type 1 constructors.
        Log("Adding constructors.");
        AddConstructorAddresses();

        // Get the battle.net.dll base address.
        Log("Getting module handle of battle.net.dll");
        HMODULE const hBattlenetDll = GetModuleHandleA("battle.net.dll");
        if (!hBattlenetDll) {
            throw Error("Could not get battle.net.dll module handle.");
        }

        // Use base address for arithmetics.
        DWORD const battlenetBase =
            reinterpret_cast<DWORD>(hBattlenetDll);

        // Rebase function.
        auto const Rebaser = [&](CtorTuple& ctor) {
            std::get<0>(ctor) += (battlenetBase - kIDA_BASE);
        };

        // Rebase constructors.
        Log("Rebasing constructor pointers.");
        std::for_each(g_messageCtors.begin(), g_messageCtors.end(), Rebaser);

        // Detour function.
        auto const MakeDetour = [&](CtorTuple& ctor) {
            // Get the constructor address.
            DWORD const address = std::get<0>(ctor);
            Log(boost::format("Detouring 0x%08x.") % address);

            // Build detour code.
            Log("Building AsmJit code.");
            AsmJit::Assembler a;
            AsmJit::Label L_getEip = a.newLabel();
            a.call(L_getEip);
            a.bind(L_getEip);
            DWORD trampolineAddr =
                reinterpret_cast<DWORD>(&CtorTrampoline);
            a.push(trampolineAddr);
            a.ret();

            // Get them bytes.
            Log("Compiling AsmJit code.");
            BYTE* detourCode = static_cast<BYTE*>(a.make());
            if (!detourCode) {
                throw Error("AsmJit::Compiler::make failed.");
            }

            // Patch memory.
            Log("Creating raw patch.");
            std::vector<BYTE> data(detourCode, detourCode + a.getCodeSize());
            auto patch = new Hades::Memory::PatchRaw(g_memory,
                reinterpret_cast<PVOID>(address), data);
            Log("Applying patch.");
            patch->Apply();

            // Keep pointer to the patch.
            Log("Adding patch to list.");
            g_patches.push_back(PatchTuple(
                std::shared_ptr<Hades::Memory::PatchRaw>(patch),
                address));
        };

        // Detour constructors.
        Log("Creating and applying constructor detours.");
        std::for_each(g_messageCtors.begin(), g_messageCtors.end(), MakeDetour);

        // Initialization done.
        Log("Completed initialization.");
    } catch (std::exception const& ex) {
        Log(boost::diagnostic_information(ex));
    } catch (...) {
        Log("Unknown error.");
    }
}

void __stdcall Log(std::string const& message) {
    if (!message.empty()) {
        // Send debug message.
        OutputDebugStringA(("*** msg-ctor-calls-log: " + message).c_str());

        // Open file.
        if (!g_logFile.is_open()) {
            g_logFile.open(kLOG_FILE_PATH, std::ofstream::app | std::ofstream::out);
        }

        // Write to file.
        g_logFile << boost::str(boost::format("[%08x] %s\n") % std::time(nullptr) % message);
        g_logFile.close();
    }
}

void __stdcall Log(boost::format const& fmt) {
    Log(boost::str(fmt));
}

void __declspec(naked) __cdecl CtorTrampoline() {
    __asm {
        push ebp
        mov ebp, esp
        push edx
        mov edx, [ebp+0x04]
        sub edx, 0x05
        push edx
        push ecx
        call CtorHandler
        pop edx
        leave
        add esp, 0x04
        retn
    }
}

DWORD __stdcall CtorHandler(DWORD messagePtr, DWORD ctorAddress) {
    // Find the patch that belongs to this constructor address.
    // TODO: Use a map instead..
    for (size_t i = 0; i < g_patches.size(); ++i) {
        if (std::get<1>(g_patches[i]) == ctorAddress) {
            // Log call.
            Log(boost::format("0x%08x: %s constructor called.")
                % ctorAddress % std::get<1>(g_messageCtors[i]));

            // Remove patch temporarily.
            std::get<0>(g_patches[i])->Remove();

            // Call original constructor.
            auto const origCtor = reinterpret_cast<DWORD (__thiscall*)(DWORD)>(
                ctorAddress);
            DWORD const result = origCtor(messagePtr);

            // Reapply patch.
            std::get<0>(g_patches[i])->Apply();

            // Return original value.
            return result;
        }
    }

    // This is bad.
    Log(boost::format("Fatal error: Unknown constructor address 0x%08x.") % ctorAddress);
    return 0;
}

void __cdecl AddConstructorAddresses() {
    // Note: This list is not complete.
    g_messageCtors.push_back(CtorTuple(0x3cc26de0, "ModuleLoadRequest"));
    g_messageCtors.push_back(CtorTuple(0x3cc27130, "ModuleLoadResponse"));
    g_messageCtors.push_back(CtorTuple(0x3cc27530, "ModuleMessageRequest"));
    g_messageCtors.push_back(CtorTuple(0x3cc27a00, "LogonRequest"));
    g_messageCtors.push_back(CtorTuple(0x3cc28360, "LogonResponse"));
    g_messageCtors.push_back(CtorTuple(0x3cc28770, "AuthModuleVariant"));
    g_messageCtors.push_back(CtorTuple(0x3cc28d10, "AuthAgreementLocale"));
    g_messageCtors.push_back(CtorTuple(0x3cc2a410, "AuthenticationConfig"));
    g_messageCtors.push_back(CtorTuple(0x3cc2af40, "AuthModuleConfig"));
    g_messageCtors.push_back(CtorTuple(0x3cc2b260, "AuthAgreement"));
    g_messageCtors.push_back(CtorTuple(0x3ccc4d80, "GetLoadRequest"));
    g_messageCtors.push_back(CtorTuple(0x3ccc4df0, "GetLoadRequest2"));
    g_messageCtors.push_back(CtorTuple(0x3ccc50b0, "ServerState"));
    g_messageCtors.push_back(CtorTuple(0x3ccc5120, "ServerState2"));
    g_messageCtors.push_back(CtorTuple(0x3ccc5610, "PoolStateRequest"));
    g_messageCtors.push_back(CtorTuple(0x3ccc5680, "PoolStateRequest2"));
    g_messageCtors.push_back(CtorTuple(0x3ccc5b80, "ServerInfo"));
    g_messageCtors.push_back(CtorTuple(0x3ccc5c00, "ServerInfo2"));
    g_messageCtors.push_back(CtorTuple(0x3ccc6460, "PoolStateResponse"));
    g_messageCtors.push_back(CtorTuple(0x3ccc64e0, "PoolStateResponse2"));
    g_messageCtors.push_back(CtorTuple(0x3ccf26b0, "NO_RESPONSE"));
    g_messageCtors.push_back(CtorTuple(0x3ccf2720, "NO_RESPONSE2"));
    g_messageCtors.push_back(CtorTuple(0x3ccf29e0, "Address"));
    g_messageCtors.push_back(CtorTuple(0x3ccf2cf0, "ProcessId"));
    g_messageCtors.push_back(CtorTuple(0x3ccf2d50, "ProcessId2"));
    g_messageCtors.push_back(CtorTuple(0x3ccf3170, "ObjectAddress"));
    g_messageCtors.push_back(CtorTuple(0x3ccf31e0, "ObjectAddress2"));
    g_messageCtors.push_back(CtorTuple(0x3ccf3340, "NoData"));
    g_messageCtors.push_back(CtorTuple(0x3ccf33b0, "NoData2"));
    g_messageCtors.push_back(CtorTuple(0x3ccf39f0, "Address2"));
    g_messageCtors.push_back(CtorTuple(0x3cd10aa0, "ConnectRequest"));
    g_messageCtors.push_back(CtorTuple(0x3cd10e20, "ConnectResponse"));
    g_messageCtors.push_back(CtorTuple(0x3cd10e90, "ConnectResponse2"));
    g_messageCtors.push_back(CtorTuple(0x3cd11230, "BoundService"));
    g_messageCtors.push_back(CtorTuple(0x3cd11290, "BoundService2"));
    g_messageCtors.push_back(CtorTuple(0x3cd11690, "BindResponse"));
    g_messageCtors.push_back(CtorTuple(0x3cd11710, "BindResponse2"));
    g_messageCtors.push_back(CtorTuple(0x3cd119e0, "EchoRequest"));
    g_messageCtors.push_back(CtorTuple(0x3cd11ea0, "EchoResponse"));
    g_messageCtors.push_back(CtorTuple(0x3cd12360, "DisconnectRequest"));
    g_messageCtors.push_back(CtorTuple(0x3cd12740, "DisconnectNotification"));
    g_messageCtors.push_back(CtorTuple(0x3cd12c60, "NullRequest"));
    g_messageCtors.push_back(CtorTuple(0x3cd12f90, "EncryptRequest"));
    g_messageCtors.push_back(CtorTuple(0x3cd148c0, "BindRequest"));
    g_messageCtors.push_back(CtorTuple(0x3cd387a0, "Var"));
    g_messageCtors.push_back(CtorTuple(0x3cd38c90, "Resource"));
    g_messageCtors.push_back(CtorTuple(0x3cd38d00, "Resource2"));
    g_messageCtors.push_back(CtorTuple(0x3cd39230, "ServerAddress"));
    g_messageCtors.push_back(CtorTuple(0x3cd392a0, "ServerAddress2"));
    g_messageCtors.push_back(CtorTuple(0x3cd396e0, "SupplementalService"));
    g_messageCtors.push_back(CtorTuple(0x3cd39740, "SupplementalService2"));
    g_messageCtors.push_back(CtorTuple(0x3cd39be0, "ServiceShard"));
    g_messageCtors.push_back(CtorTuple(0x3cd39c50, "ServiceShard2"));
    g_messageCtors.push_back(CtorTuple(0x3cd3ab20, "Var2"));
    g_messageCtors.push_back(CtorTuple(0x3cd3b130, "ProgramResources"));
    g_messageCtors.push_back(CtorTuple(0x3cd3b1b0, "ProgramResources2"));
    g_messageCtors.push_back(CtorTuple(0x3cd3b450, "ServerSet"));
    g_messageCtors.push_back(CtorTuple(0x3cd3b4e0, "ServerSet2"));
    g_messageCtors.push_back(CtorTuple(0x3cd3ba30, "ListenPoint"));
    g_messageCtors.push_back(CtorTuple(0x3cd3bab0, "ListenPoint2"));
    g_messageCtors.push_back(CtorTuple(0x3cd3bd00, "ServiceConfig"));
    g_messageCtors.push_back(CtorTuple(0x3cd3bd80, "ServiceConfig2"));
    g_messageCtors.push_back(CtorTuple(0x3cd3c260, "RPCServerConfig"));
    g_messageCtors.push_back(CtorTuple(0x3cd3c2e0, "RPCServerConfig2"));
    g_messageCtors.push_back(CtorTuple(0x3cd3c590, "ProcessConfig"));
    g_messageCtors.push_back(CtorTuple(0x3cd3c670, "ProcessConfig2"));
}