#include <stdio.h>
#include <winsock.h>

#include "uodef.h"

int SendBuffer(SOCKET s, void *Buffer, int Length)
{
  // First :
  //   test if data can be send
  //   1) if an error occurs : abort!
  //   2) if data can be send : break the loop
  while(1)
  {
    FD_SET totest = {1, {s}};
    TIMEVAL timeout = {0, 10}; // 0 seconds, 10 ms
    int rv = select(0, NULL, &totest, NULL, &timeout);
    if(rv == SOCKET_ERROR)
    {
      printf("Unable to test for send : %u bytes lost\n", Length);
      return(-2);
    }
    if(rv != 0)
      break;
  }

  if(send(s, (char *) Buffer, Length, 0) != Length)
  {
    printf("Unable to send buffer : %u bytes lost\n", NULL, MB_OK);
    return(-1);
  }

  return(0);
}

int SendBufferEncrypted(SOCKET s, void *Buffer, int Length)
{
  return(SendBuffer(s, Buffer, Length));
}

inline int SendLong(SOCKET s, unsigned long Long)
{
  return(SendBuffer(s, &Long, sizeof(Long)));
}

void main(int argc, char *argv[])
{
  WSADATA wsaData;
  const WORD VersionRequired = MAKEWORD(1, 1);

  // My favorite test defaults...
  unsigned short serverport = 2593;
  char *servername = "novusopiate.com";
  char *username = "admin";
  char *password = "banme";

  // Check the parameter
  if(argc < 4)
  {
    printf("%s\n", "Parameter error\nUsage : exe <username> <password> <servername> [port]\nPort defaults to 2593\n\nPress ENTER/RETURN to exit...");
    getchar();
    return;
  }

  // Configure the program depending on the parameters
  username = argv[1];
  password = argv[2];
  servername = argv[3];
  serverport = argc > 4 ? atoi(argv[4]) : 0;
  if(serverport == 0)
    serverport = 2593;

  // Initialize the socket library
  if(WSAStartup(VersionRequired, &wsaData ) != 0)
  {
    printf("%s\n", "Socket init error!");
    return;
  }

  // Verify version
  if(wsaData.wVersion != VersionRequired)
  {
    WSACleanup();

    printf("%s\n", "Socket version error!");
    return;
  }

  // Convert the servername to an IP
  unsigned long serverip = inet_addr(servername);
  if(serverip == INADDR_NONE)
  {
    HOSTENT *hostp = gethostbyname(servername);
    if(hostp != NULL)
      memcpy(&serverip, hostp->h_addr, sizeof(serverip));
  }

  // Prepare the initial socket
  SOCKET ServerSocket = socket(AF_INET, SOCK_STREAM, 0);
  SOCKADDR_IN ServerAddress = {AF_INET, htons(serverport), {127, 0, 0, 1}};
  if(serverip != INADDR_NONE)
    ServerAddress.sin_addr.S_un.S_addr = serverip;

  // Connect to the server
  if(connect(ServerSocket, (SOCKADDR *) &ServerAddress, sizeof(ServerAddress)) != 0)
  {
    printf("Unable to connect to <%s>.\n", servername);
    return;
  }
  else
    printf("We're connected to <%s:%u>.\nTrying to log on with %s, %s...\n\n", servername, serverport, username, password);

  // Send a login start (a 0)
  // Note :
  //   - uox will disconnect if you send 0x78563412
  //   - to act as a xgm client you must send 0xFFFFFFFF
  //   - to connect to OSI you must send 0x0100007F and use encryption : untested
  SendLong(ServerSocket, 0x00000000);

  // Send the request to get the server list
  // NOTE :
  //  - username and password is required for OSI/UOX(/Wolfpack?)
  //  - sphere(/tus) do not require a username/password
  //  - POL : untested
  UO_GetServerList GetServerList = {0x80, "", "", 0xFF};
  strncpy(GetServerList.userid  , username, sizeof(GetServerList.userid));
  strncpy(GetServerList.password, password, sizeof(GetServerList.password));
  SendBuffer(ServerSocket, &GetServerList, sizeof(GetServerList));

  bool SendPassword = false;
  bool UseEncryption = false;
  long Encryptor = 0x00000000;
  while(1)
  {
    if(SendPassword)
    {
      // Reset the flag so we do not resend the password
      SendPassword = false;

      // Send the login (userid+password)
      UO_ClientLogin ClientLogin = {0x91, 0, "", ""};
      strncpy(ClientLogin.userid  , username, sizeof(ClientLogin.userid));
      strncpy(ClientLogin.password, password, sizeof(ClientLogin.password));
      SendBuffer(ServerSocket, &ClientLogin, sizeof(ClientLogin));

      // Start using encryption from now on
      UseEncryption = true;
    }

    // First :
    //   test if there is data to read
    //   1) if an error occurs : break the loop
    //   2) if no data is ready to read : restart the loop
    FD_SET totest = {1, {ServerSocket}};
    TIMEVAL timeout = {0, 10}; // 0 seconds, 10 ms
    int rv = select(0, &totest, NULL, NULL, &timeout);
    if(rv == SOCKET_ERROR)
      break;
    if(rv == 0)
      continue;

    // Second :
    //   get the number of bytes to read
    //   1) if an error occurs : break the loop
    //   2) if no data is ready to read (-> connection reset) : break the loop
    //   3) maximum number of bytes to read at once is 32767 (very unlikely though)
    //      (but we check for sure, we do not wish to get exploited)
    long ByteCount;
    if(ioctlsocket(ServerSocket, FIONREAD, (u_long *) &ByteCount) != 0)
      break;
    if(ByteCount <= 0)
      break;
    if(ByteCount > 32767)
      ByteCount = 32767;

    // Peek at the code (aka command) of the packet
    // (we do this to determine the packet size)
    UO_Code Code;
    if(recv(ServerSocket, (char *) &Code, sizeof(Code), MSG_PEEK) != sizeof(Code))
    {
      printf("Unable to peek at code : %u bytes lost!\n", sizeof(Code));
      return;
    }

    // If encryption is enabled we should decrypt the packet first
    if(UseEncryption)
    {
      // We are not interested in decryption, just the initial byte for status report
      // NOTE : if login fails we should decrypt the message to know the actual reason
      if(Code.code == 0x81)
      {
        printf("%s\n", "Login OK (Sphere/Tus Server)");
        break;
      }
      if(Code.code == 0x59)
      {
        printf("%s\n", "Login failure (Sphere/Tus Server)");
        break;
      }
      printf("%s\n", "Login/Server Type unknown :(");
      break;
    }
    else
    {
      if(Code.code == 0xA9)
      {
        // UOX Server Login OK
        printf("%s\n", "Login OK (UOX/Wolfpack with Ignition)");
        break;
      }
      if(Code.code == 0x82)
      {
        // Login Failure
        printf("%s\n", "Login failure (UOX/Wolfpack with Ignition)");
        break;
      }
    }
    
    // Unknown packet? -> msg this to the user and exit
    if(UO_StructureSize[Code.code] == 0)
    {
      char Buffer[32];
      sprintf(Buffer, "%d bytes", ByteCount);

      // Assume that the bad packet is as long as the data length in the socket buffer
      // Allocate memory to receive this (bad) data
      char *Buffer4 = (char *) malloc(ByteCount);
      if(Buffer4 != NULL)
      {
        // Receive the data
        if(recv(ServerSocket, Buffer4, ByteCount, 0) != ByteCount)
        {
          printf("Unable to receive unknown packet type! : %u bytes lost\n", ByteCount);
          return;
        }

        // Display the data to the user as a hex string
        char *Buffer2 = (char *) malloc(ByteCount * 3 + 1);
        if(Buffer2 != NULL)
        {
          Buffer2[0] = '\0';
          for(int i = 0; i < ByteCount; i ++)
          {
            char Buffer3[2 + 1];
            sprintf(Buffer3, "%02X ", (unsigned char) Buffer4[i]);
            strcat(Buffer2, Buffer3);
          }
          printf("%s\n%s\n", Buffer, Buffer2);

          // Release used memory
          free(Buffer2);
        }

        // Release user memory
        free(Buffer4);
      }

      break;
    }

    // If the packet isn't completely received : keep looping
    if(ByteCount < UO_StructureSize[Code.code])
      continue;

    // Handle structures that can contain extra bytes!!
    int ExtraByteCount;
    if(Code.code == 0xA8)
    {
      // Peek at the full structure
      // (so we can determine the extra byte count)
      UO_ServerList ResponseTest;    
      if(recv(ServerSocket, (char *) &ResponseTest, sizeof(ResponseTest), MSG_PEEK) != sizeof(ResponseTest))
      {
        printf("Unable to peek at extra data! : %u bytes lost\n", sizeof(ResponseTest));
        return;
      }

      // Determine the number of extra bytes to receive
      ExtraByteCount = ntohs(ResponseTest.servercount) * sizeof(UO_ServerInfo);
    }
    else
    {
      // The packet doesn't have extra bytes (lucky us)
      ExtraByteCount = 0;
    }

    // Keep looping if not all the extra bytes are received
    if((UO_StructureSize[Code.code] + ExtraByteCount) > ByteCount)
      continue;

    // Only get this packet
    // if the number of bytes to receive is greater than the current packet size
    // -> keep them for the next loop
    ByteCount = UO_StructureSize[Code.code];

    // Receive the packet (as a C union data type)
    UO_CodeData Response;
    if(recv(ServerSocket, (char *) &Response, ByteCount, 0) != ByteCount)
    {
      printf("Unable to receive the packet! : %u bytes lost\n", ByteCount);
      return;
    }

    // Receive the extra bytes (if any)
    // (and allocate seperate memory for this extra data)
    void *ExtraResponse = ExtraByteCount > 0 ? malloc(ExtraByteCount) : NULL;
    if(ExtraResponse != NULL)
      if(recv(ServerSocket, (char *) ExtraResponse, ExtraByteCount, 0) != ExtraByteCount)
      {
        printf("Unable to receive the extra data! : %u bytes lost\n", ExtraByteCount);
        return;
      }

    // Process the packet
    if(Response.code == 0xA8) // Server List
    {
      #define ServerInfo ((UO_ServerInfo *) ExtraResponse)[i]
      for(int i = 0; i < ntohs(Response.serverlist.servercount); i ++)
      {
        char Buffer[512], ServerCount[64];
        sprintf(ServerCount, "server info (%d servers)", ntohs(Response.serverlist.servercount));
        sprintf(Buffer, "server id = %d (%d%% full)\n     server name = %.30s\n     server id(/ip) = %d.%d.%d.%d\n", ntohs(ServerInfo.id), ServerInfo.percentfull, ServerInfo.name, (ServerInfo.ipaddress >> 0) & 0xFF, (ServerInfo.ipaddress >> 8) & 0xFF, (ServerInfo.ipaddress >> 16) & 0xFF, (ServerInfo.ipaddress >> 24) & 0xFF);
        printf("%02u : %s", i, Buffer);
      }
      #undef ServerInfo

      // Now redirect to the default server
      if(ntohs(Response.serverlist.servercount) > 0)
      {
        UO_RedirectClientRequest RedirectClientRequest = {0xA0, ((UO_ServerInfo *) ExtraResponse)[0].id};
        SendBuffer(ServerSocket, &RedirectClientRequest, sizeof(RedirectClientRequest));
      }
    }
    else if(Response.code == 0x8C) // Relay
    {
      // Close the socket 
      // NOTE : you might wanna keep this socket open on an OSI-server
      closesocket(ServerSocket);

      // Initialize the new socket
      ServerSocket = socket(AF_INET, SOCK_STREAM, 0);
      SOCKADDR_IN ServerAddress = {AF_INET, Response.redirectclient.port, {127, 0, 0, 1}};
      ServerAddress.sin_addr.S_un.S_addr = Response.redirectclient.ipaddress;

      // Open a new socket to the relayed IP
      if(connect(ServerSocket, (SOCKADDR *) &ServerAddress, sizeof(ServerAddress)) != 0)
      {
        printf("Relay failed!\n");
        return;
      }

      // Send the encryptor
      // NOTE : this 4-byte number is used to modify the packet encryption
      //        we send a 0, the OSI client sends your internal/external IP
      SendLong(ServerSocket, Encryptor);

      // Set the flag to send the password
      SendPassword = true;
    }
    else 
    {
      // Doh!, packet not processed : display the data as an hex string instead
      // (note : extra data is not shown)
      char *Buffer2 = (char *) malloc(ByteCount * 3 + 1);
      if(Buffer2 != NULL)
      {
        Buffer2[0] = '\0';
        for(int i = 0; i < ByteCount; i ++)
        {
          char Buffer3[2 + 1];
          sprintf(Buffer3, "%02X ", ((unsigned char *) &Response)[i]);
          strcat(Buffer2, Buffer3);
        }

        printf("Known Packet not implemented :\n%s\n", Buffer2);

        // Release the temporary used memory
        free(Buffer2);
      }

      break;
    }

    // Release the extra memory (if any)
    if(ExtraResponse != NULL)
      free(ExtraResponse);
  }

  // Clean up...
  WSACleanup();

  // Wait for exit
  printf("\n%s", "Press ENTER/RETURN to exit...");
  getchar();;
}