#include<stdio.h>
#include<stdlib.h>
#include<stdbool.h>
#include<string.h>
#include<pcap/pcap.h>
#include<stdbool.h>
#include<inttypes.h>
#include<netinet/ether.h>
#include<arpa/inet.h>
#include"linkedList.h"
#include"timecap.h"

#ifndef PROG_NAME
#define PROG_NAME "a.out"
#endif

// Seconds
#define KEY_EPS 1

#define SSH_SETUP_PACKET_COUNT 4

#define FLAG_SHORT_HELP "-h"
#define FLAG_LONG_HELP  "--help"

#define FLAG_SHORT_DIFF "-t"
#define FLAG_LONG_DIFF  "--time-diff"

#define FLAG_SHORT_LENG "-l"
#define FLAG_LONG_LENG  "--len-diff"

#define FLAG_SHORT_VERB "-v"
#define FLAG_LONG_VERB  "--verbose"

/* Format of keylog lines in log file.
 * 4 = the fd SSH reads user input from (don't know why it isn't 0 but it isn't)
 * 16384 = How many bytes SSH reads as a time
 * The last %c is there to check that the entire string matched correctly.
 * This does not read the entire line.  Call read_line after using this.
 */
#define KEYLOG_FORMAT "%lf read(%*d, \"%m[^\"]\", 16384) %c"

#define S_TO_NS 1000000000
#define S_TO_US 1000000

#define DEFAULT_SYN_E {3, 0}

#define DEFAULT_TIM_E 20

typedef struct {
  FILE* keylog;
  char* keylogName;
  struct timeval keylogStart;
  List* timecaps;
  struct timeval synEpsiolon;
  size_t keyEpsilon;
  bool verbose;
  size_t line_count;
} Args;

void check_error(int cond, char* str) {
  if (cond) {
    perror(str);
    exit(1);
  }
}

void print_usage(int exitCode) {
  printf("Usage: %s [options] keylog pcap [pcap ...]\n"
         "Pair a keylog to a SSH TCP stream.\n\n"
         "Options:\n"
         "%-20s Show this help message and exit\n"
         "%-20s Print reasons for rejecting pcap files\n"
         "%-20s The allowable difference in time (in seconds) between keylog \n"
         "%-20s start and flow start\n"
         "%-20s The allowable difference in estimated keys between keylog and\n"
         "%-20s flow \n",
         PROG_NAME,
         "  "FLAG_SHORT_HELP", "FLAG_LONG_HELP,
         "  "FLAG_SHORT_VERB", "FLAG_LONG_VERB,
         "  "FLAG_SHORT_DIFF", "FLAG_LONG_DIFF, "",
         "  "FLAG_SHORT_LENG", "FLAG_LONG_LENG, "");
  exit(exitCode);
}

struct timeval double_to_timeval(double in) {
  struct timeval time;
  time.tv_sec = in;
  time.tv_usec = (in - time.tv_sec) * S_TO_US;
  return time;
}

void parse_flag(char** argv, Args* res, size_t* i) {
  double d;
  char* arg = argv[*i];
  if (strcmp(arg, FLAG_SHORT_HELP) == 0 || strcmp(arg, FLAG_LONG_HELP) == 0) {
    print_usage(0);
  }
  else if (strcmp(arg, FLAG_SHORT_VERB) == 0 ||
           strcmp(arg, FLAG_LONG_VERB) == 0) {
    res->verbose = true;
  }
  else if (strcmp(arg, FLAG_SHORT_DIFF) == 0 ||
           strcmp(arg, FLAG_LONG_DIFF) == 0) {
    if (sscanf(argv[++*i], " %lf", &d) == 1 && d > 0) {
      res->synEpsiolon = double_to_timeval(d);
    }
    else {
      print_usage(1);
    }
  }
  else if (0 == strcmp(arg, FLAG_SHORT_LENG) ||
           0 == strcmp(arg, FLAG_LONG_LENG)) {
    if (1 != sscanf(argv[++*i], " %zu", &res->keyEpsilon)) {
      print_usage(1);
    }
  }
}

/* Assumes each pcap has at least one packet and each pcap doesn't overlap.
 */
static void add_timecap(List* pcaps, char* filename) {
  char* errBuf  = alloca(PCAP_ERRBUF_SIZE + strlen(filename) +
                         sizeof(NO_PACKET_ERR));
  TimeCap* cap = new_timecap(filename, errBuf);
  if (cap == NULL) {
    fprintf(stderr, "%s: %s\n", filename, errBuf);
  }
  else {
    list_add_head(pcaps, cap);
  }
}

static int read_line(FILE* keylog) {
  int c;
  while ('\n' != (c = fgetc(keylog)) && EOF != c) {
    continue;
  }
  return EOF == c? EOF : 0;
}

static struct timeval get_keylog_start(FILE* keylog) {
  double seconds;
  if (fscanf(keylog, " %lf", &seconds) != 1) {
    fprintf(stderr, "Failed to read timestame from keylog file.\n");
    exit(1);
  }
  read_line(keylog);
  return double_to_timeval(seconds);
}

size_t count_lines(FILE* file) {
  size_t res = 1;
  while (EOF != read_line(file)) {
    ++res;
  }
  return res;
}

static Args get_args(int argc, char** argv) {
  Args res = {NULL, NULL, {0, 0}, new_list(NULL), DEFAULT_SYN_E, DEFAULT_TIM_E};
  size_t i = 1;
  while (argv[i]) {
    if (argv[i][0] == '-') {
      parse_flag(argv, &res, &i);
    }
    else if (res.keylog) {
      add_timecap(res.timecaps, argv[i]);
    }
    else {
      res.keylog = fopen(argv[i], "r");
      res.keylogName = argv[i];
      check_error(res.keylog == NULL, argv[i]);
      res.keylogStart = get_keylog_start(res.keylog);
      res.line_count = count_lines(res.keylog);
    }
    ++i;
  }
  if (res.keylog == NULL || res.timecaps == NULL) {
    print_usage(1);
  }
  return res;
}

void timeval_subtract(struct timeval *result,
                      struct timeval *x,
                      struct timeval *y) {
  *result = *x;
  if (result->tv_usec < y->tv_usec) {
    result->tv_sec -= 1;
    result->tv_usec += 1*S_TO_US;
  }
  result->tv_sec -= y->tv_sec;
  result->tv_usec -= y->tv_usec;
  result->tv_sec = abs(result->tv_sec);
  result->tv_usec = abs(result->tv_usec);
}

int timercmp_eps(struct timeval* a, struct timeval* b, struct timeval eps) {
  struct timeval diff;
  timeval_subtract(&diff, a, b);
  return timercmp(&diff, &eps, <);
}

#define LOG(ARGS, FMT, ...) do {                \
  if ((ARGS)->verbose) {                        \
    fprintf(stderr, (FMT), ##__VA_ARGS__);      \
  }                                             \
} while (0)

int is_match(void* tc, void* a) {
  TimeCap* cap = (TimeCap*) tc;
  Args* args = (Args*) a;
  int res;
  if (timercmp_eps(&args->keylogStart, &cap->time, args->synEpsiolon)) {
    res = args->line_count <= cap->pack_count + args->keyEpsilon &&
      args->line_count >= cap->pack_count - args->keyEpsilon;
    if (!res) {
      LOG(args, "PCAP %s had incompatible line numbers.\n", cap->name);
      LOG(args, "%zu keys, %zu packets.\n", args->line_count, cap->pack_count);
    }
    return res;
  }
  else {
    LOG(args, "PCAP %s had incompatible timestamp.\n", cap->name);
    return 0;
  }
}

void cleanup(Args args) {
  delete_list(args.timecaps, delete_timecap);
  fclose(args.keylog);
}

int main(int argc, char** argv) {
  Args args = get_args(argc, argv);
  TimeCap* cap = (TimeCap*) list_find(args.timecaps, &args, is_match);
  int exitStatus = 0;
  if (cap != NULL) {
    printf("%s\n", cap->name);
  }
  else {
    fprintf(stderr, "Failed to find matching flow file for %s.\n",
            args.keylogName);
    exitStatus = 1;
  }
  cleanup(args);
  return exitStatus;
}