#include "check_csgo.h"
#include <vector>

using namespace std;

char *servicename = (char*)"CS:GO";

void printVersion() {
  cout << "check_csgo v" << VERSION << endl << endl;
}

void printHelp(bool longVersion) {
  if(longVersion) {
    printVersion();
    cout << "Check CS:GO DS instance." << endl << endl;
    printHelp(false);
    cout << "Options:" << endl;
    cout << " -h" << endl;
    cout << "    Print detailed help screen" << endl;
    cout << " -V" << endl;
    cout << "    Print version information" << endl;
    cout << " -H HOSTADDRESS" << endl;
    cout << "    Host where the Source DS is running" << endl;
    cout << " -p" << endl;
    cout << "    Port where the Source DS is listening. Default is 27015." << endl << endl;
    return;
  }
  cout << "Usage: " << endl << "check_csgo [-hV] -H HOSTADDRESS [-p PORT]" << endl << endl;
}

SOURCEDS_PACKET* getSourceDSResponse(int s, sockaddr_in *server) {
  char *data = new char[DATAGRAM_LENGTH];
  int dataLength = recvMsg(s, data, DATAGRAM_LENGTH, server);
  if(dataLength == -1) {
    return NULL;
  }
  char *header = (char*) "\xFF\xFF\xFF\xFF";
  if(memcmp(data,header,HDR_SIZE)) {
    //Invalid response, or non-supported multi-packet response format
    return NULL;
  }
  SOURCEDS_PACKET *response = new SOURCEDS_PACKET;
  response->data = data+HDR_SIZE;
  response->length = dataLength;
  return response;
}

int check_csgo(char *hostname, uint16_t port, SERVERINFO *server_info) {
  struct sockaddr_in si;
  struct hostent *host = NULL;
  int s;
  int timeout = 10;
  setupSocket(port,hostname,host,timeout,&si,&s);
  string discard;

  std::vector<uint8_t> info_query;
  info_query.resize(25);
  memcpy(info_query.data(), "\xFF\xFF\xFF\xFF\x54Source Engine Query\x00", 25);

  sendMsg(s, reinterpret_cast<char*>(info_query.data()), info_query.size(), &si);
  SOURCEDS_PACKET *response = getSourceDSResponse(s,&si);
  if(response == NULL) {
    return 2;
  }

  if(response->data[0] == 'A') {
    // Challenge response mechanism
    uint32_t challenge;
    size_t current_pkt_size = info_query.size();
    info_query.resize(current_pkt_size + 4);
    memcpy(info_query.data() + current_pkt_size, response->data+1, 4);
    delete[] (response->data-HDR_SIZE);
    delete response;

    sendMsg(s, reinterpret_cast<char*>(info_query.data()), info_query.size(), &si);
    response = getSourceDSResponse(s, &si);
    if(response == NULL) {
      return 2;
    }
  }

  ssize_t offset = 2;
  server_info->name = string(response->data+offset,strlen(response->data+offset));
  offset += server_info->name.size()+1;
  server_info->map = string(response->data+offset,strlen(response->data+offset));
  offset += server_info->map.size()+1;
  discard = string(response->data+offset,strlen(response->data+offset));
  offset += discard.size() + 1;
  server_info->game = string(response->data+offset,strlen(response->data+offset));
  offset += server_info->game.size()+1;
  offset += 2;
  server_info->players = (uint8_t) *((response->data)+offset);
  server_info->max_players = (uint8_t) *((response->data)+(++offset));

  delete[] (response->data-HDR_SIZE);
  delete response;

  if(server_info->players == server_info->max_players) {
    return 1;
  }

  return 0;
}

int main(int argc, char **argv) {
  struct itimerval timer;
  timer.it_value.tv_sec = 10;
  timer.it_value.tv_usec = 0;
  timer.it_interval.tv_sec = 0;
  timer.it_interval.tv_usec = 0;
  setitimer (ITIMER_VIRTUAL, &timer, 0);

  struct sigaction sa;
  memset (&sa, 0, sizeof (sa));
  sa.sa_handler = &timer_handler;
  sigaction (SIGVTALRM, &sa, 0);

  uint16_t port = 27015;
  char *hostname = NULL;
  int c;

  while ((c = getopt (argc, argv, "H:p:Vh")) != -1) {
    switch(c) {
      case 'H':
        hostname = optarg;
        break;
      case 'p':
        port = (uint16_t) str2int(optarg);
       	break;
      case 'V':
        printVersion();
        return 0;
      case 'h':
        printHelp(true);
        return 0;
      case '?':
        printHelp(false);
        return 3;
    }
  }

  if(hostname == NULL) {
    cout << "No HOSTADDRESS specified. Exiting." << endl;
    return 3;
  }

  SERVERINFO server_info;
  int returnCode = check_csgo(hostname,port,&server_info);

  cout << servicename;
  switch(returnCode) {
    case 0:
      cout << " OK";
      cout << " - " << hostname << " " << server_info.name << " " << server_info.game << " " << server_info.map
           << " " << (int)server_info.players << "/" << (int)server_info.max_players << " players" << endl ;
      break;

    case 1:
      cout << " WARNING - Server is full";
      break;
    case 2:
      cout << " CRITICAL - No response";
      break;
  }
  return returnCode;
}