PT2Hook - dllmain

// dllmain.cpp : Defines the entry point for the DLL application.
#include "stdafx.h"
#include "Logging.h"



int(WINAPI* pWSASend)(SOCKET s, LPWSABUF lpBuffers, DWORD dwBufferCount, LPDWORD lpNumbersOfBytesSent, DWORD dwFlags, LPWSAOVERLAPPED lpOverlapped, LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine) = WSASend;
int (WINAPI *pWSARecv)(SOCKET s, LPWSABUF lpBuffers, DWORD dwBufferCount, LPDWORD lpNumberOfBytesRecvd, LPDWORD lpFlags, LPWSAOVERLAPPED lpOverlapped, LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine) = WSARecv;
void (CALLBACK *pOriginalCompletion_Recv)(DWORD, DWORD, LPWSAOVERLAPPED, DWORD) = NULL;
void (CALLBACK *pOriginalCompletion_Send)(DWORD, DWORD, LPWSAOVERLAPPED, DWORD) = NULL;


static bool CompletionHook_Recv = false;
static bool CompletionHook_Send = false;
static WSABUF bufferPointer;
                                      //   opcode hour minute seconds
const char* szRecvdPacketDumpFormat = "%s\\S2C_%02X_%02d_%02d_%02d.dat";
const char* szSentPacketDumpFormat = "%s\\C2S_%02X_%02d_%02d_%02d.dat";
const char* szPacketsDumpPath = "C:\\PT2_Packets";

enum PacketType
{
    CLIENT_TO_SERVER = 2,
    SERVER_TO_CLIENT = 4,
};

void WINAPI PrintPacket(DWORD length, PacketType mode)
{
    unsigned short packetOpcode = *(unsigned short*)(bufferPointer.buf);

    // generate fileName
    char szPacketFileName[64] = "";

    Logging::UpdateTime();

    switch (mode)
    {
    case PacketType::CLIENT_TO_SERVER:
        _snprintf(szPacketFileName, 64, szSentPacketDumpFormat,
            szPacketsDumpPath,
            packetOpcode,
            Logging::currentDate->tm_hour,
            Logging::currentDate->tm_min,
            Logging::currentDate->tm_sec
            );
        Logging::PrintDebug("Dumping Client => Server packet with length %d. - Possible Opcode: %02X. \n", length, packetOpcode);
        break;

    case PacketType::SERVER_TO_CLIENT:
        _snprintf(szPacketFileName, 64, szRecvdPacketDumpFormat,
            szPacketsDumpPath,
            packetOpcode,
            Logging::currentDate->tm_hour,
            Logging::currentDate->tm_min,
            Logging::currentDate->tm_sec
            );
        Logging::PrintDebug("Dumping Server => Client packet with length %d. - Possible Opcode: %02X. \n", length, packetOpcode);
        break;
    }

    // create file and write content
    FILE* pFile = fopen(szPacketFileName, "wb+");
    fwrite(bufferPointer.buf, 1, length, pFile);
    fclose(pFile);

}

// MyCompletionCallback
void CALLBACK MyCompletionCallback_Recv(DWORD dwError, DWORD dwTransferred, LPWSAOVERLAPPED lpOverlapped, DWORD dwFlags)
{

    //ogging::PrintDebug("Recv Completion Routine called!\n");
    DWORD bytesReceived = 0;

    if (dwTransferred == 0)
    {
        Logging::PrintDebug("dwTransferred = 0!\n");
    }
    else if (dwTransferred > 0)
    {
        bytesReceived = dwTransferred;
        PrintPacket(bytesReceived, PacketType::SERVER_TO_CLIENT);
    }
    pOriginalCompletion_Recv(dwError, bytesReceived, lpOverlapped, dwFlags);
}

void CALLBACK MyCompletionCallback_Send(DWORD dwError, DWORD dwTransferred, LPWSAOVERLAPPED lpOverlapped, DWORD dwFlags)
{
    //ogging::PrintDebug("Send Completion Routine called!\n");
    DWORD bytesReceived = 0;

    if (dwTransferred == 0)
    {
        Logging::PrintDebug("dwTransferred = 0!\n");
    }
    else if (dwTransferred > 0)
    {
        bytesReceived = dwTransferred;
        PrintPacket(bytesReceived, PacketType::CLIENT_TO_SERVER);
    }
    pOriginalCompletion_Send(dwError, bytesReceived, lpOverlapped, dwFlags);
}

int WINAPI MyWSASend(SOCKET s, LPWSABUF lpBuffers, DWORD dwBufferCount, LPDWORD lpNumberOfBytesSent, DWORD lpFlags, LPWSAOVERLAPPED lpOverlapped, LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine)
{
    if (!CompletionHook_Send)
    {
        //Logging::PrintDebug("First WSARecv called, trying to hook Completion Routine...\n");
        // set original function pointer
        pOriginalCompletion_Send = lpCompletionRoutine;

        DetourUpdateThread(GetCurrentThread());
        DetourTransactionBegin();
        DetourAttach(&reinterpret_cast<PVOID&>(pOriginalCompletion_Send), MyCompletionCallback_Send);
        if (DetourTransactionCommit() == NO_ERROR)
        {
            //Logging::PrintDebug("Completion Routine for WSARecv hooked!\n");
            CompletionHook_Send = true;
        }
    }
    return pWSASend(s, lpBuffers, dwBufferCount, lpNumberOfBytesSent, lpFlags, lpOverlapped, lpCompletionRoutine);
}
int WINAPI MyWSARecv(SOCKET s, LPWSABUF lpBuffers, DWORD dwBufferCount, LPDWORD lpNumberOfBytesRecvd, LPDWORD lpFlags, LPWSAOVERLAPPED lpOverlapped, LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine)
{
    if (!CompletionHook_Recv)
    {
        //Logging::PrintDebug("First WSARecv called, trying to hook Completion Routine...\n");
        // set original function pointer
        pOriginalCompletion_Recv = lpCompletionRoutine;

        DetourUpdateThread(GetCurrentThread());
        DetourTransactionBegin();
        DetourAttach(&reinterpret_cast<PVOID&>(pOriginalCompletion_Recv), MyCompletionCallback_Recv);
        if (DetourTransactionCommit() == NO_ERROR)
        {
            //Logging::PrintDebug("Completion Routine for WSARecv hooked!\n");
            CompletionHook_Recv = true;
        }
    }

    //Logging::PrintDebug("WSARecv called! - Updating the bufferPointer variable.\n");

    // copy buffer data
    CopyMemory(&bufferPointer, lpBuffers, sizeof(WSABUF));

    return pWSARecv(s, lpBuffers, dwBufferCount, lpNumberOfBytesRecvd, lpFlags, lpOverlapped, lpCompletionRoutine);
}

void BeginOperations(HMODULE me)
{
    // alloc console.
    AllocConsole();
    freopen("CONOUT$", "w", stdout);
    SetConsoleTitle("PT2 Sniffer - by Sheen");

    // start log
    Logging::Start();

    // check packet directory
    DWORD dwAttributes = GetFileAttributes(szPacketsDumpPath);
    BOOL dirExists = (dwAttributes == INVALID_FILE_ATTRIBUTES) ? false : dwAttributes & FILE_ATTRIBUTE_DIRECTORY;

    if (!dirExists)
    {
        BOOL success = CreateDirectory(szPacketsDumpPath, NULL);
        if (!success)
        {
            Logging::PrintDebug("Failed to create directory %s. Packet dump function will not work properly.", szPacketsDumpPath);
        }
        else
        {
            Logging::PrintDebug("Directory %s created to save the packet dumps.", szPacketsDumpPath);
        }
    }


    DisableThreadLibraryCalls(me);
    //Logging::PrintDebug("Hooking WSARecv...\n");

    DetourRestoreAfterWith();

    DetourTransactionBegin();
    DetourUpdateThread(GetCurrentThread());

    DetourAttach(&reinterpret_cast<PVOID&>(pWSARecv), MyWSARecv);
    DetourAttach(&reinterpret_cast<PVOID&>(pWSASend), MyWSASend);

    if (DetourTransactionCommit() == NO_ERROR)
        Logging::PrintDebug("Hooks OK.\n");
}

void EndOperations(HMODULE me)
{
    DetourTransactionBegin();
    DetourUpdateThread(GetCurrentThread());

    DetourDetach(&reinterpret_cast<PVOID&>(pWSARecv), MyWSARecv);
    DetourDetach(&reinterpret_cast<PVOID&>(pWSASend), MyWSASend);
    DetourAttach(&reinterpret_cast<PVOID&>(pOriginalCompletion_Recv), MyCompletionCallback_Recv);
    DetourAttach(&reinterpret_cast<PVOID&>(pOriginalCompletion_Send), MyCompletionCallback_Send);

    DetourTransactionCommit();
}

BOOL APIENTRY DllMain( HMODULE hModule,
                       DWORD  ul_reason_for_call,
                       LPVOID lpReserved
                     )
{
    switch (ul_reason_for_call)
    {
    case DLL_PROCESS_ATTACH:

        BeginOperations(hModule);

        break;
    case DLL_THREAD_ATTACH:
    case DLL_THREAD_DETACH:

        EndOperations(hModule);
        break;
    case DLL_PROCESS_DETACH:
        break;
    }
    return TRUE;
}