/* @(#)client.c
 */

#include <unistd.h>

#define _GNU_SOURCE
#include <getopt.h>
#include <string.h>

#include <sys/socket.h>
#include <netinet/in.h>
#include <netinet/ip.h> /* superset of previous */
#include <sys/socket.h>
#include <stdio.h>

#include <gcrypt.h>

#define QUIT(sd) write(sd, "QUIT\n",5)

const unsigned short port = 1234;
#define FLAG_SIZE 16

static void
hexdump(unsigned char *buf, unsigned int len)
{     
        while (len--)  
                printf("%02x ", *buf++);
        printf("\n");
}

static void flag_to_ascii(const unsigned char *bin, unsigned char *ascii){
  int i;
  memset(ascii,0,32);
  for (i=0; i<FLAG_SIZE;i++) 
    sprintf(ascii+(2*i), "%02x", bin[i]);
}



int calc_hmac(const unsigned char *flag, unsigned char *digest){
  char *mdigest;
  gcry_md_hd_t mcontext;

/*   gcry_md_init(); */
  gcry_control (GCRYCTL_DISABLE_SECMEM);
  gcry_md_open(&mcontext, GCRY_MD_MD5, GCRY_MD_FLAG_HMAC);
  gcry_md_setkey(mcontext, "CTF secret key <g>", 18);
  gcry_md_write(mcontext, flag, strlen(flag));
  mdigest = gcry_md_read(mcontext,0);
  
  memcpy(digest, mdigest, FLAG_SIZE);

  gcry_md_close(mcontext);
  
  return 0;
}


int connect_to_server(const char *ip, unsigned short port) {
  struct sockaddr_in serv;
  int sd,r ;

  serv.sin_family = AF_INET;
  serv.sin_port = htons(port);

  if ((sd = socket(PF_INET, SOCK_STREAM, 0)) == -1) {return -3;}
  if (inet_pton(AF_INET, ip, &(serv.sin_addr)) <= 0) { return -2;}
  
  r = connect(sd, (struct sockaddr*)&serv, (socklen_t)sizeof(serv));
  if (r) return r;
  return sd;
}


int get_flag(int sd, int l, char *buf, int len){
  FILE *f;
  char s[20];
  
  sprintf(s, "GET %d\n", l);
  f = fdopen(sd, "r+");
  memset(buf,0,len);
  fputs(s,f);
  fgets(buf,len,f);
  if (strcmp(buf,"No module!\n") == 0) return 0;
  return 2;
}

inline static int verify_flag(unsigned char *ascii1, unsigned char *ascii2, int len) {
  return (memcmp(ascii1, ascii2, len)==0)?2:1;
}



int main(int argc, char **argv) {
  /* 
     return values:
     0: remote malfunction | service not up | ...
     1: service up, but no flag
     2: success
     3: internal error
  */
  int ret = 2;
  int c, sd = -1;
  unsigned char *host = NULL, *flag=NULL;
  int bget,bset,len;
  
  bget = bset = 0;
  len = 32;
  
  while (1) {
    c = getopt(argc, argv, "gf:h:l:s");
    if (c == -1 ) break;
    switch (c) {
    case 'h':
      host = strdup(optarg);
      break;
    case 'l':
      len = atoi(optarg);
      break;
    case 'f':
      flag = strdup(optarg);
      break;
    case 'g':
      if (bset) {
	printf("You may either GET or SET, but not both!\n");
	ret = 3;
	goto out;
      }
      bget = 1;
      break;
    case 's':
      if (bget) {
	printf("You may either GET or SET, but not both!\n");
	ret = 3;
	goto out;
      }
      bset = 1;
      break;
    } /* switch */
    
  }

  printf ("len: %d\n", len);

  sd = connect_to_server(host, port);
  if (sd < -1) {
    ret = 3;
    goto out;
  } else if (sd  == -1) {
    ret = 0;
    goto out;
  }
  
  if (bget){
    /* get FLAG */
    unsigned char current_flag[abs(len)+1], hmac[16], vrfy_flag[abs(len)+1];
    
    ret = get_flag(sd, len, current_flag, abs(len)+1);
    if (ret != 2) {
      QUIT(sd);
      goto out;
    }
    calc_hmac(flag, hmac);
    flag_to_ascii(hmac, vrfy_flag);
    ret = verify_flag(current_flag, vrfy_flag, abs(len));
    QUIT(sd);
    goto out;

  } else if (bset) {

    /* set FLAG */
    unsigned char current_flag[abs(len)+1], hmac[16], vrfy_flag[abs(len)+1], tmp[256];
    FILE *f;

    ret = get_flag(sd, len, current_flag, abs(len)+1);
    if (ret != 2) {
      QUIT(sd);
      goto out;
    }
    calc_hmac(current_flag, hmac);
    flag_to_ascii(hmac, vrfy_flag);
    sprintf(tmp, "PUT %s %s\r\n", vrfy_flag, flag);
    write(sd, tmp, strlen(tmp));
    QUIT(sd);

  } else {

    printf ("What do you want?\n");
    ret = 3;
    goto out;
  }
        
  out:
  if (host) free(host);
  if (flag) free(flag);
  if (sd >= 0) close(sd);
   
  return ret;
}

