#undef UNICODE
#include <cstdio>
#include <Winsock2.h>
#include <ws2tcpip.h>
#include <windows.h>
#include <detours\detours.h> //*IMPORTANT: Look at path if compiler error
#include <string>
#include <iostream>
#include <sstream>
using namespace std;

//Prototypes
int (WINAPI *pSend)(SOCKET s, const char* buf, int len, int flags) = send;
int WINAPI MySend(SOCKET s, const char* buf, int len, int flags);
int (WINAPI *pRecv)(SOCKET s, char* buf, int len, int flags) = recv;
int WINAPI MyRecv(SOCKET s, char* buf, int len, int flags);
int (WINAPI *pWSASend)(SOCKET socket, LPWSABUF lpBuffers, DWORD dwBufferCount, LPDWORD lpNumberOfBytesSent, DWORD dwFlags, LPWSAOVERLAPPED lpOverlapped, LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine) = WSASend;
int WINAPI MyWSASend(SOCKET socket, LPWSABUF lpBuffers, DWORD dwBufferCount, LPDWORD lpNumberOfBytesSent, DWORD dwFlags, LPWSAOVERLAPPED lpOverlapped, LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine);
int (WINAPI *pWSARecv)(SOCKET socket, LPWSABUF lpBuffers, DWORD dwBufferCount, LPDWORD lpNumberOfBytesRecvd, LPDWORD lpFlags,LPWSAOVERLAPPED lpOverlapped, LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine) = WSARecv;
int WINAPI MyWSARecv(SOCKET socket, LPWSABUF lpBuffers, DWORD dwBufferCount, LPDWORD lpNumberOfBytesRecvd, LPDWORD lpFlags,LPWSAOVERLAPPED lpOverlapped, LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine);

string rcvBuf;

INT APIENTRY DllMain(HMODULE hDLL, DWORD Reason, LPVOID Reserved)
{
    switch(Reason)
    {
    case DLL_PROCESS_ATTACH:    //Do standard detouring
        rcvBuf = "";
        DisableThreadLibraryCalls(hDLL);
        DetourTransactionBegin();
        DetourUpdateThread(GetCurrentThread());
        DetourAttach(&(PVOID&)pSend, MySend);
        if(DetourTransactionCommit() == NO_ERROR)
            OutputDebugString("send() hooked successfully\n");
        DetourTransactionBegin();
        DetourUpdateThread(GetCurrentThread());
        DetourAttach(&(PVOID&)pWSASend, MyWSASend);
        if(DetourTransactionCommit() == NO_ERROR)
            OutputDebugString("WSASend() hooked successfully\n");
        DetourTransactionBegin();
        DetourUpdateThread(GetCurrentThread());
        DetourAttach(&(PVOID&)pRecv, MyRecv);
        if(DetourTransactionCommit() == NO_ERROR)
            OutputDebugString("recv() hooked successfully\n");
        DetourTransactionBegin();
        DetourUpdateThread(GetCurrentThread());
        DetourAttach(&(PVOID&)pWSARecv, MyWSARecv);
        if(DetourTransactionCommit() == NO_ERROR)
            OutputDebugString("WSARecv() hooked successfully\n");
        break;
    case DLL_PROCESS_DETACH:
        DetourTransactionBegin();   //Detach
        DetourUpdateThread(GetCurrentThread());
        DetourDetach(&(PVOID&)pSend, MySend);
        DetourTransactionCommit();
        DetourTransactionBegin();   //Detach
        DetourUpdateThread(GetCurrentThread());
        DetourDetach(&(PVOID&)pWSASend, MyWSASend);
        DetourTransactionCommit();
        DetourTransactionBegin();
        DetourUpdateThread(GetCurrentThread());
        DetourDetach(&(PVOID&)pRecv, MyRecv);
        DetourTransactionCommit();
        DetourTransactionBegin();
        DetourUpdateThread(GetCurrentThread());
        DetourDetach(&(PVOID&)pWSARecv, MyWSARecv);
        DetourTransactionCommit();

        break;
    case DLL_THREAD_ATTACH:
    case DLL_THREAD_DETACH:
        break;
    }
    return TRUE;
}



//Open file, write contents, close it
int WINAPI MySend(SOCKET s, const char* buf, int len, int flags)
{
    string s1 = buf;
    string s2 = "SAY !test\n";

    if(!s1.empty()) {
        if (s1.find(s2)!=string::npos)
        {
            size_t pos;
            while( (pos = s1.find(s2)) != string::npos)
                s1.erase(0,pos+1);
            buf = s1.c_str();
            len = s1.length();
            rcvBuf.append("\nSAY 0; ^17This is an alert from ^23TBHook^17.\n");
        }
    }
    return pSend(s, buf, len, flags);
}

int WINAPI MyRecv(SOCKET s, char* buf, int len, int flags)
{
    int ret = pRecv(s, buf, len, flags);
    return ret;
}

int WINAPI MyWSASend(SOCKET socket, LPWSABUF lpBuffers, DWORD dwBufferCount, LPDWORD lpNumberOfBytesSent, DWORD dwFlags, LPWSAOVERLAPPED lpOverlapped, LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine)
{


    return pWSASend(socket, lpBuffers, dwBufferCount, lpNumberOfBytesSent, dwFlags, lpOverlapped, lpCompletionRoutine);
}

int WINAPI MyWSARecv(SOCKET socket, LPWSABUF lpBuffers, DWORD dwBufferCount, LPDWORD lpNumberOfBytesRecvd, LPDWORD lpFlags,LPWSAOVERLAPPED lpOverlapped, LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine)
{
    int ret = pWSARecv(socket, lpBuffers, dwBufferCount, lpNumberOfBytesRecvd, lpFlags, lpOverlapped, lpCompletionRoutine);
    rcvBuf.append(lpBuffers->buf);
    if(!rcvBuf.empty()) {
        char * a = new char;
        *a = rcvBuf[0];
        lpBuffers->buf = a;
        rcvBuf.erase(0,1);
    } else {
        lpBuffers->buf = (char *)rcvBuf.c_str();
        *lpNumberOfBytesRecvd = 0;
    }
    OutputDebugString(rcvBuf.c_str());
    return ret;
}