#include <stdio.h>
#include <string.h>
#define INCL_WINSOCK_API_PROTOTYPES 0
#define INCL_WINSOCK_API_TYPEDEFS 1
#define WINSOCK_API_LINKAGE
#include <winsock2.h>
#include <shlobj.h>
#include <ws2spi.h>

// hack to exclude down-level inline code, which depends on ws2_32
// calls, which are not available due to INCL_WINSOCK_API_PROTOTYPES
#define _WSPIAPI_H_
#include <ws2tcpip.h>
#undef _WSPIAPI_H_

#include <sporder.h>

// code assumes 32-bit pointers
#if ((!defined(_WIN32) && !defined(WIN32)) || (defined(_WIN64) || defined(WIN64)))
  #error x64 platform is not supported!
#endif

//-- Definitions --//

// including <btypes.h>
#ifndef __cplusplus
  #ifndef __bool_true_false_are_defined
    typedef char bool;
    #define false ((bool)(1==0))
    #define true  ((bool)(1==1))
    #define __bool_true_false_are_defined
  #endif
#endif

#ifndef _WCHAR_T_DEFINED
  typedef unsigned short wchar_t;
  #define _WCHAR_T_DEFINED
  #define WCHAR_MIN 0
  #define WCHAR_MAX ((wchar_t)-1)
#endif  /* _WCHAR_T_DEFINED */

#define BUFP(p,pos) (((unsigned char*)(p)) + (pos))

#ifndef _ARRAYSIZE
  #define _ARRAYSIZE(x) (sizeof(x)/sizeof((x)[0]))
#endif  /* _ARRAYSIZE */
// end of <btypes.h>

// including <sock.h>
#define SOCK_JOINIP(a,b,c,d) ( ( ( (a)&0xFF ) << 24 ) | ( ( (b)&0xFF ) << 16 ) | ( ( (c)&0xFF ) << 8 ) | ( (d)&0xFF ) )
#define SOCK_SPLITIP(ip) ( ( (ip) >> 24 )&0xFF ), ( ( (ip) >> 16 )&0xFF ), ( ( (ip) >> 8 )&0xFF ), ( (ip)&0xFF )
// end if <sock.h>


//-- Function Typedefs --//
typedef void (__stdcall* LPFNSETDLLDIRECTORYA)(LPCSTR lpPathName);
typedef int (__stdcall* LPFNWSAFDISSET)(SOCKET fd, fd_set* set);

//-- Globals --//
static bool use_dll = TRUE;
static HANDLE fHandle = INVALID_HANDLE_VALUE;
static HINSTANCE l_hRealDll = NULL;
static HINSTANCE l_hFakeDll = NULL;

static struct WSFUNCS
{
  //By Ordinal
  LPFN_ACCEPT accept;
  LPFN_BIND bind;
  LPFN_CLOSESOCKET closesocket;
  LPFN_CONNECT connect;
  LPFN_GETPEERNAME getpeername;
  LPFN_GETSOCKNAME getsockname;
  LPFN_GETSOCKOPT getsockopt;
  LPFN_HTONL htonl;
  LPFN_HTONS htons;
  LPFN_IOCTLSOCKET ioctlsocket;
  LPFN_INET_ADDR inet_addr;
  LPFN_INET_NTOA inet_ntoa;
  LPFN_LISTEN listen;
  LPFN_NTOHL ntohl;
  LPFN_NTOHS ntohs;
  LPFN_RECV recv;
  LPFN_RECVFROM recvfrom;
  LPFN_SELECT select;
  LPFN_SEND send;
  LPFN_SENDTO sendto;
  LPFN_SETSOCKOPT setsockopt;
  LPFN_SHUTDOWN shutdown;
  LPFN_SOCKET socket;
  LPFN_GETHOSTBYADDR gethostbyaddr;
  LPFN_GETHOSTBYNAME gethostbyname;
  LPFN_GETSERVBYNAME getservbyname;
  LPFN_GETHOSTNAME gethostname;
  LPFN_WSAGETLASTERROR WSAGetLastError;
  LPFN_WSASTARTUP WSAStartup;
  LPFN_WSACLEANUP WSACleanup;
  LPFNWSAFDISSET __WSAFDIsSet;
  
  //By Name
  LPFN_WSASEND WSASend;
  LPFN_WSARECV WSARecv;
  LPFN_WSAGETOVERLAPPEDRESULT WSAGetOverlappedResult;
}
l_WsFuncs = { 0 };

//-- Helper Functions --//
void Dump2File(SOCKET s, char* buf, int len, int ncount, const char* prefix)
{
  if (use_dll && ncount && ncount != SOCKET_ERROR && fHandle != INVALID_HANDLE_VALUE)
  {
    //Print the bytes to a string in hex format
	int dummy, numBytes = strlen(prefix) + 2 + len*2 + 2;
    char* szBuffer = (char*)malloc(numBytes);
	
    wsprintfA(szBuffer, "%s\r\n", prefix); // strlen(prefix) + 2
	
    for (int i = 0; i < len; i++)
	{
      wsprintfA(szBuffer + i*2 + 4, "%02X", buf[i]); // len*2
    }
	
    wsprintfA(szBuffer + len*2 + 4 , "\r\n"); // 2
	
	WriteFile(fHandle, szBuffer, numBytes, &dummy, NULL);
  }
}

//-- Exported Functions --//
SOCKET __stdcall accept(SOCKET s, struct sockaddr* addr, int* addrlen)
{
  return l_WsFuncs.accept(s, addr, addrlen);
}

int __stdcall bind(SOCKET s, const struct sockaddr* name, int namelen)
{
  return l_WsFuncs.bind(s, name, namelen);
}

int __stdcall closesocket(SOCKET s)
{
  return l_WsFuncs.closesocket(s);
}

int __stdcall connect(SOCKET s, const struct sockaddr* name, int namelen)
{
  return l_WsFuncs.connect(s, name, namelen);
}

int __stdcall getpeername(SOCKET s, struct sockaddr* name, int* namelen)
{
  return l_WsFuncs.getpeername(s, name, namelen);
}

int __stdcall getsockname(SOCKET s, struct sockaddr* name, int* namelen)
{
  return l_WsFuncs.getsockname(s, name, namelen);
}

int __stdcall getsockopt(SOCKET s, int level, int optname, char *optval, int *optlen)
{
  return l_WsFuncs.getsockopt(s, level, optname, optval, optlen);
}

u_long __stdcall htonl(u_long hostlong)
{
  return l_WsFuncs.htonl(hostlong);
}

u_short __stdcall htons(u_short hostshort)
{
  return l_WsFuncs.htons(hostshort);
}

int __stdcall ioctlsocket(SOCKET s, long cmd, u_long* argp)
{
  return l_WsFuncs.ioctlsocket(s, cmd, argp);
}

unsigned long __stdcall inet_addr(const char* cp)
{
  return l_WsFuncs.inet_addr(cp);
}

char* __stdcall inet_ntoa(struct in_addr in)
{
  return l_WsFuncs.inet_ntoa(in);
}

int __stdcall listen(SOCKET s, int backlog)
{
  return l_WsFuncs.listen(s, backlog);
}

u_long __stdcall ntohl(u_long netlong)
{
  return l_WsFuncs.ntohl(netlong);
}

u_short __stdcall ntohs(u_short netshort)
{
  return l_WsFuncs.ntohs(netshort);
}

int __stdcall recv(SOCKET s, char* buf, int len, int flags)
{
  //Recv first
  int n = l_WsFuncs.recv(s, buf, len, flags);
  
  //Dump the data to our file
  Dump2File(s, buf, len, n, "RR");
  
  //Return result of recv function
  return n;
}

int __stdcall recvfrom(SOCKET s, char *buf, int len, int flags, struct sockaddr *from, int *fromlen)
{
  return l_WsFuncs.recvfrom(s, buf, len, flags, from, fromlen);
}

int __stdcall select(int nfds, fd_set* readfds, fd_set* writefds, fd_set*exceptfds, const struct timeval* timeout)
{
  return l_WsFuncs.select(nfds, readfds, writefds, exceptfds, timeout);
}

int __stdcall send(SOCKET s, const char* buf, int len, int flags)
{
  //Send First
  int n = l_WsFuncs.send(s, buf, len, flags);
  
  //Dump the data to our file
  Dump2File(s, buf, len, n, "SS");
  
  //Return result of send function
  return n;
}

int __stdcall sendto(SOCKET s, const char* buf, int len, int flags, const struct sockaddr* to, int tolen)
{
  return l_WsFuncs.sendto(s, buf, len, flags, to, tolen);
}

int __stdcall setsockopt(SOCKET s, int level, int optname, const char* optval, int optlen)
{
  return l_WsFuncs.setsockopt(s, level, optname, optval, optlen);
}

int __stdcall shutdown(SOCKET s, int how)
{
  return l_WsFuncs.shutdown(s, how);
}

SOCKET __stdcall socket(int af, int type, int protocol)
{
  return l_WsFuncs.socket(af, type, protocol);
}

struct hostent* __stdcall gethostbyaddr(const char* addr, int len, int type)
{
  return l_WsFuncs.gethostbyaddr(addr, len, type);
}

struct hostent* __stdcall gethostbyname(const char* name)
{
  return l_WsFuncs.gethostbyname(name);
}

struct servent* __stdcall getservbyname(const char* name, const char* proto)
{
  return l_WsFuncs.getservbyname(name, proto);
}

int __stdcall gethostname(char* name, int namelen)
{
  return l_WsFuncs.gethostname(name, namelen);
}

int __stdcall WSAStartup(WORD wVersionRequested, LPWSADATA lpWSAData)
{
  return l_WsFuncs.WSAStartup(wVersionRequested, lpWSAData);
}

int __stdcall WSACleanup(void)
{
  return l_WsFuncs.WSACleanup();
}

int __stdcall WSAGetLastError(void)
{
  return l_WsFuncs.WSAGetLastError();
}

int __stdcall __WSAFDIsSet(SOCKET fd, fd_set* set)
{
  return l_WsFuncs.__WSAFDIsSet(fd, set);
}

__declspec(dllexport) int __stdcall WSASend(SOCKET s, LPWSABUF lpBuffers, DWORD dwBufferCount, LPDWORD lpNumberOfBytesSent, DWORD dwFlags, LPWSAOVERLAPPED lpOverlapped, LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine)
{
  return l_WsFuncs.WSASend(s, lpBuffers, dwBufferCount, lpNumberOfBytesSent, dwFlags, lpOverlapped, lpCompletionRoutine);
}

_declspec(dllexport) int __stdcall WSARecv(SOCKET s, LPWSABUF lpBuffers, DWORD dwBufferCount, LPDWORD lpNumberOfBytesRecvd, LPDWORD lpFlags, LPWSAOVERLAPPED lpOverlapped, LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine)
{
  return l_WsFuncs.WSARecv(s, lpBuffers, dwBufferCount, lpNumberOfBytesRecvd, lpFlags, lpOverlapped, lpCompletionRoutine);
}

_declspec(dllexport) BOOL __stdcall WSAGetOverlappedResult(SOCKET s, LPWSAOVERLAPPED lpOverlapped, LPDWORD lpcbTransfer, BOOL fWait, LPDWORD lpdwFlags)
{
  return l_WsFuncs.WSAGetOverlappedResult(s, lpOverlapped, lpcbTransfer, fWait, lpdwFlags);
}

//-- Entry Point (Setup everything & Unload everything in this) --//
BOOL CALLBACK DllMain(HINSTANCE hInstance, DWORD dwReason, LPVOID lpReserved)
{
  char szSysPath[MAX_PATH], szLibPath[MAX_PATH];  
  switch(dwReason)
  {
    case DLL_PROCESS_ATTACH:
	{
      l_hFakeDll = hInstance;
	
      //Check whether we will pursue dumping packets with User
      if (MessageBoxA(NULL, "Dump packets?", "Packet Dumper", MB_YESNO|MB_ICONQUESTION)==IDNO)
      {
        use_dll = FALSE;
      }
	
      if (use_dll)
      {
	    //Create a Handle to our dump file (name is calculated from current timestamp
	    SYSTEMTIME st;
	    GetLocalTime(&st);
	    char fName[32];//Enough space to fit the name in all cases
	    wsprintfA(fName, "./PacketDump_%4d%2d%2d_%2d%2d%2d.txt", st.wYear, st.wMonth, st.wDay, st.wHour, st.wMinute, st.wSecond);
		
        fHandle = CreateFile( fName, GENERIC_WRITE, FILE_SHARE_READ, NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
	    if (fHandle == INVALID_HANDLE_VALUE)
	    {
	      MessageBoxA(NULL, "Dump file not created!", "File Error", MB_OK|MB_ICONSTOP);
		  use_dll = FALSE;
	    }
      }
	
      {// Rewrite system DLL load path
        LPFNSETDLLDIRECTORYA lpfnSetDllDirectoryA = (LPFNSETDLLDIRECTORYA)GetProcAddress(GetModuleHandle("KERNEL32.DLL"), "SetDllDirectoryA");
      
        if(lpfnSetDllDirectoryA)
        {
          lpfnSetDllDirectoryA(".");
        }
      }
      
	  //Load the Original dll and get the Real Function addresses.
      GetSystemDirectoryA(szSysPath, _ARRAYSIZE(szSysPath));
      wsprintfA(szLibPath, "%s\\ws2_32.dll", szSysPath);
    
      l_hRealDll = LoadLibraryA(szLibPath);    
      if(l_hRealDll)
      {
        if(l_hRealDll==l_hFakeDll)
        {// oops
          FreeLibrary(l_hRealDll);
          l_hRealDll = NULL;
          MessageBoxA(NULL, "Attempted self-load.", "Library Error", MB_OK|MB_ICONSTOP);
          return FALSE;
        }
		
#define LOADFUNC(x,y) l_WsFuncs.y = (x)GetProcAddress(l_hRealDll, #y)
        LOADFUNC(LPFN_ACCEPT,accept);
        LOADFUNC(LPFN_BIND,bind);
        LOADFUNC(LPFN_CLOSESOCKET,closesocket);
        LOADFUNC(LPFN_CONNECT,connect);
        LOADFUNC(LPFN_GETPEERNAME,getpeername);
        LOADFUNC(LPFN_GETSOCKNAME,getsockname);
        LOADFUNC(LPFN_GETSOCKOPT,getsockopt);
        LOADFUNC(LPFN_HTONL,htonl);
        LOADFUNC(LPFN_HTONS,htons);
        LOADFUNC(LPFN_IOCTLSOCKET,ioctlsocket);
        LOADFUNC(LPFN_INET_ADDR,inet_addr);
        LOADFUNC(LPFN_INET_NTOA,inet_ntoa);
        LOADFUNC(LPFN_LISTEN,listen);
        LOADFUNC(LPFN_NTOHL,ntohl);
        LOADFUNC(LPFN_NTOHS,ntohs);
        LOADFUNC(LPFN_RECV,recv);
        LOADFUNC(LPFN_RECVFROM,recvfrom);
        LOADFUNC(LPFN_SELECT,select);
        LOADFUNC(LPFN_SEND,send);
        LOADFUNC(LPFN_SENDTO,sendto);
        LOADFUNC(LPFN_SETSOCKOPT,setsockopt);
        LOADFUNC(LPFN_SHUTDOWN,shutdown);
        LOADFUNC(LPFN_SOCKET,socket);
        LOADFUNC(LPFN_GETHOSTBYADDR,gethostbyaddr);
        LOADFUNC(LPFN_GETHOSTBYNAME,gethostbyname);
        LOADFUNC(LPFN_GETHOSTNAME,gethostname);
        LOADFUNC(LPFN_GETSERVBYNAME,getservbyname);
        LOADFUNC(LPFN_WSAGETLASTERROR,WSAGetLastError);
        LOADFUNC(LPFN_WSASTARTUP,WSAStartup);
        LOADFUNC(LPFN_WSACLEANUP,WSACleanup);
        LOADFUNC(LPFNWSAFDISSET, __WSAFDIsSet);
        
        LOADFUNC(LPFN_WSASEND, WSASend);
        LOADFUNC(LPFN_WSARECV, WSARecv);
        LOADFUNC(LPFN_WSAGETOVERLAPPEDRESULT, WSAGetOverlappedResult);
#undef LOADFUNC
      
      }
      else
      {
        return FALSE;
      }
      break;
	}
	
    case DLL_PROCESS_DETACH:
	{
	  //Close the dump file if open
	  if (fHandle != INVALID_HANDLE_VALUE)
	  {
	    CloseHandle(fHandle);
	    fHandle = INVALID_HANDLE_VALUE;
	  }
	  
	  //Cleanup the Real DLL loaded.
      if (l_hRealDll)
      {
       
	    l_WsFuncs.WSACleanup();
        FreeLibrary(l_hRealDll);
        l_hRealDll = NULL;
      }
	  
      l_hFakeDll = NULL;
      break;
	}
  }
  return TRUE;
  
  // unused
  lpReserved;
}