#include "packet.h"
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <err.h>

#include <poll.h>
#include <sys/wait.h>
#include <unistd.h>
#include <stdbool.h>

struct proc {
   pid_t pid;
   int fds[2];
};

static void
close_fd(int *fd)
{
   if (*fd >= 0)
      close(*fd);
}

static void
proc_close(struct proc *proc)
{
   if (proc->pid)
      waitpid(proc->pid, NULL, 0);

   close_fd(&proc->fds[0]);
   close_fd(&proc->fds[1]);
   *proc = (struct proc){0};
}

static bool
proc_open(const char *file, char *const argv[], struct proc *out_proc)
{
   *out_proc = (struct proc){0};

   int pipes[4];
   if (pipe(&pipes[0]) != 0 /* parent */ || pipe(&pipes[2]) != 0 /* child */) {
      proc_close(out_proc);
      return false;
   }

   if ((out_proc->pid = fork()) > 0) {
      out_proc->fds[0] = pipes[3];
      out_proc->fds[1] = pipes[0];
      close(pipes[1]);
      close(pipes[2]);
      return true;
   } else {
      close(pipes[0]);
      close(pipes[3]);
      dup2(pipes[2], 0);
      dup2(pipes[1], 1);
      close(pipes[2]);
      close(pipes[1]);
      execvp(file, argv);
      _exit(0);
   }

   out_proc->fds[0] = pipes[3];
   out_proc->fds[1] = pipes[0];
   close(pipes[1]);
   close(pipes[2]);
   return true;
}

static size_t
safe_read(int fd, void *dst, const size_t dst_size)
{
   size_t off = 0;
   ssize_t ret = 0;
   for (; off < dst_size && (ret = read(fd, (char*)dst + off, dst_size - off)) > 0; off += ret);
   return (ret < 0 ? 0 : off);
}

static void
dump_packet(const uint8_t *bytes, const size_t sz, const char *prefix, const uint16_t type)
{
   const char *local = getenv("TCPLOCALPORT");

   char path[4096];
   snprintf(path, sizeof(path), "packets/%s-%s-0x%.4x.raw", local, prefix, type);

   FILE *f;
   if (!(f = fopen(path, "wb")))
      err(EXIT_FAILURE, "fopen");

   if (fwrite(bytes, 1, sz, f) != sz)
      err(EXIT_FAILURE, "fwrite");

   fclose(f);
}

static void
handle_select_character_res(union packet *packet, bool prologue)
{
   packet->hdr.type = (prologue ? PACKET_CHARACTER_ENTER_PROLOGUE : PACKET_SELECT_CHARACTER_RES);
   struct pbuf pbuf = { .packet = packet, .cursor = sizeof(packet->hdr) };
   pbuf.cursor += sizeof(uint32_t); // char id
   for (size_t i = 0; i < 32; ++i) pbuf_write(&pbuf, (uint8_t[]){0xff}, 1); // unknown
   const char *server_ip = getenv("SERVER_IP");
   server_ip = (server_ip ? server_ip : "127.0.0.1");
   pbuf_write_str(&pbuf, server_ip); // server ip
   pbuf_write(&pbuf, (uint16_t[]){10100}, sizeof(uint16_t)); // port
   for (size_t i = 0; i < 38; ++i) pbuf_write(&pbuf, (uint8_t[]){0x00}, 1); // unknown
   pbuf_flush(&pbuf);
}

static void
handle_character_enter_prologue(union packet *packet)
{
   handle_select_character_res(packet, true);
}

static void
handle_server_list(union packet *packet)
{
   struct pbuf pbuf = { .packet = packet, .cursor = sizeof(packet->hdr) };
   pbuf_write(&pbuf, (uint8_t[]){0x00, 0x01, 0x00, 0x00}, 4); // unknown
   pbuf_write(&pbuf, (uint16_t[]){10000}, sizeof(uint16_t)); // port
   pbuf_write_str(&pbuf, "benis"); // server name
   const char *server_ip = getenv("SERVER_IP");
   server_ip = (server_ip ? server_ip : "127.0.0.1");
   pbuf_write_str(&pbuf, server_ip); // server ip
   pbuf_write(&pbuf, (uint8_t[]){0x01, 0x00, 0x00, 0x00, 0x9f, 0x01, 0x00, 0x00, 0x00}, 9);
   pbuf_flush(&pbuf);
}

static bool
intercept(union packet *packet)
{
   packet_crypt(packet);

   bool ret = true;
   switch (packet->hdr.type) {
      case PACKET_SERVER_LIST:
         handle_server_list(packet);
         break;
      case PACKET_SELECT_CHARACTER_RES:
         handle_select_character_res(packet, false);
         break;
      case PACKET_CHARACTER_ENTER_PROLOGUE:
         handle_character_enter_prologue(packet);
         break;
      default:
         ret = false;
         break;
   }

   packet_crypt(packet);
   return ret;
}

static uint16_t
decrypted_packet_type(union packet *packet)
{
    const uint16_t size = packet->hdr.size;
    packet->hdr.size = sizeof(packet->hdr);
    packet_crypt(packet);
    const uint16_t type = packet->hdr.type;
    packet_crypt(packet);
    packet->hdr.size = size;
    return type;
}

int
main(int argc, char *const argv[])
{
   for (int i = 1; i < argc; ++i) {
      if (!strcmp(argv[i], "--help") || !strcmp(argv[i], "-h"))
         errx(EXIT_SUCCESS, "usage: %s cmd [args]", argv[0]);
   }

   struct proc server = {0};
   if (argc >= 2 && !proc_open(argv[1], argv + 1, &server))
      errx(EXIT_FAILURE, "failed to execute: %s", argv[1]);

   setvbuf(stdout, NULL, _IONBF, 0);

   struct pollfd pfd[] = {
      { .fd = STDIN_FILENO, .events = POLLIN },
      { .fd = server.fds[1], .events = POLLIN },
   };

   union packet packet;
   while (poll(pfd, 1 + !!server.pid, -1) > 0) {
      if (pfd[0].revents & POLLIN) {
         if (fread(packet.buf, 1, sizeof(packet.hdr), stdin) < sizeof(packet.hdr))
            errx(EXIT_SUCCESS, "client close");

         if (!packet_verify(&packet)) {
            for (size_t i = 0; i < 7; ++i) fprintf(stderr, "%.2x ", packet.buf[i]);
            errx(EXIT_FAILURE, "invalid packet");
         }

         if (fread(packet.buf + sizeof(packet.hdr), 1, packet.hdr.size - sizeof(packet.hdr), stdin) != packet.hdr.size - sizeof(packet.hdr))
            errx(EXIT_SUCCESS, "failed to read packet from client");

         const uint16_t type = decrypted_packet_type(&packet);
         dump_packet(packet.buf, packet.hdr.size, "client", type);

         if (intercept(&packet)) {
            dump_packet(packet.buf, packet.hdr.size, "client-intercepted", type);
            warnx("intercepted client packet 0x%.4x", type);
         } else {
            warnx("client packet 0x%.4x (%u)",type, packet.hdr.size);
         }

         if (server.pid && write(server.fds[0], packet.buf, packet.hdr.size) != packet.hdr.size)
            err(EXIT_FAILURE, "failed to write to server");
      } else if (pfd[1].revents & POLLIN) {
         if (read(server.fds[1], packet.buf, sizeof(packet.hdr)) < (ssize_t)sizeof(packet.hdr))
            errx(EXIT_SUCCESS, "server close");

         if (!packet_verify(&packet)) {
            for (size_t i = 0; i < 7; ++i) fprintf(stderr, "%.2x ", packet.buf[i]);
            errx(EXIT_FAILURE, "invalid packet");
         }

         if (safe_read(server.fds[1], packet.buf + sizeof(packet.hdr), packet.hdr.size - sizeof(packet.hdr)) != (packet.hdr.size - sizeof(packet.hdr)))
            errx(EXIT_SUCCESS, "failed to read packet from server");

         const uint16_t type = decrypted_packet_type(&packet);
         dump_packet(packet.buf, packet.hdr.size, "server", type);

         if (intercept(&packet)) {
            dump_packet(packet.buf, packet.hdr.size, "server-intercepted", type);
            warnx("intercepted server packet 0x%.4x",type);
         } else {
            warnx("server packet 0x%.4x (%u)", type, packet.hdr.size);
         }

         if (fwrite(packet.buf, 1, packet.hdr.size, stdout) != packet.hdr.size)
            err(EXIT_FAILURE, "failed to write to client");
      }
   }

   proc_close(&server);
   return EXIT_SUCCESS;
}