#include <errno.h>
#include <fcntl.h>
#include <libwifi.h>
#include <linux/if_packet.h>
#include <net/if.h>
#include <netinet/ether.h>
#include <netinet/ip.h>
#include <pcap/pcap.h>
#include <poll.h>
#include <pthread.h>
#include <stdio.h>
#include <sys/ioctl.h>
#include <unistd.h>

#include "addr_list.h"
#include "beacon_data.h"

static unsigned char kBroadcastAddress[] = "\xFF\xFF\xFF\xFF\xFF\xFF";

static pthread_mutex_t socket_mutex = PTHREAD_MUTEX_INITIALIZER;

static pthread_mutex_t seq_mutex = PTHREAD_MUTEX_INITIALIZER;
static int seq_number = 0;

int get_interface_mac_address(int sock, const char *interface, unsigned char *output)
{
  struct ifreq ifr = {0};
  strncpy(ifr.ifr_name, interface, IFNAMSIZ - 1);

  if (ioctl(sock, SIOCGIFHWADDR, &ifr) < 0)
  {
    return -1;
  }

  memcpy(output, ifr.ifr_hwaddr.sa_data, 6);
  return 0;
}

int get_interface_index(int sock, const char *interface, int *index)
{
  struct ifreq ifr = {0};
  strncpy(ifr.ifr_name, interface, IFNAMSIZ - 1);

  if (ioctl(sock, SIOCGIFINDEX, &ifr) < 0)
  {
    return -1;
  }

  *index = ifr.ifr_ifindex;
  return 0;
}

int get_next_seq_number() {
  pthread_mutex_lock(&seq_mutex);
  int next = seq_number++;
  pthread_mutex_unlock(&seq_mutex);

  return next;
}

int make_beacon_frame(int index, const unsigned char tx_addr[6], int seq_number, unsigned char **buf)
{
  struct libwifi_beacon beacon = {0};

  if (libwifi_create_beacon(&beacon, kBroadcastAddress, tx_addr, tx_addr, NULL, 13))
  {
    printf("Could not create beacon frame.\n");
    return -3;
  }

  beacon.frame_header.seq_control.sequence_number = seq_number;
  beacon.fixed_parameters.capabilities_information = (1 << CAPABILITIES_ESS) | (1 << CAPABILITIES_SHORT_PREAMBLE) | (1 << CAPABILITIES_SHORT_SLOT);
  beacon.fixed_parameters.beacon_interval = 200;

  static const unsigned char supported_rates[] = {0x82, 0x84, 0x0b, 0x16, 0x24, 0x30, 0x48, 0x6c};
  if (libwifi_quick_add_tag(&beacon.tags, TAG_SUPP_RATES, supported_rates, 8))
  {
    printf("Could not add supported rates tag.\n");
    return -7;
  }

  static const unsigned char tim_bitmap[] = {0x01, 0x02, 0x00, 0x00};
  if (libwifi_quick_add_tag(&beacon.tags, TAG_TIM, tim_bitmap, 4)) {
    printf("Could not add TIM tag.\n");
    return -7;
  }

  static const unsigned char nothing = 0;
  libwifi_quick_add_tag(&beacon.tags, TAG_ERP, &nothing, 1);
  libwifi_quick_add_tag(&beacon.tags, 47, &nothing, 1);

  static const unsigned char extended_supported_rates[] = {0x0c, 0x12, 0x18, 0x60};
  if (libwifi_quick_add_tag(&beacon.tags, TAG_EXTENDED_SUPPORTED_RATES, extended_supported_rates, 4))
  {
    printf("Could not add extended supported rates tag.\n");
    return -7;
  }

  unsigned char payload_data[BEACON_PAYLOAD_LENGTH + 8];
  payload_data[0] = 0x00;
  payload_data[1] = 0x09;
  payload_data[2] = 0xBF;
  payload_data[3] = 0x00;
  payload_data[4] = 0xFF;
  payload_data[5] = 0xFF;
  payload_data[6] = 0x00;
  payload_data[7] = 0x00;
  memcpy(payload_data + 8, kBeaconPayloads[index], BEACON_PAYLOAD_LENGTH);

  if (libwifi_quick_add_tag(&beacon.tags, TAG_VENDOR_SPECIFIC, payload_data, BEACON_PAYLOAD_LENGTH + 8))
  {
    printf("Could not add beacon data tag.\n");
    return -6;
  }

  size_t beacon_size = libwifi_get_beacon_length(&beacon);
  unsigned char *beacon_output = (unsigned char *)malloc(beacon_size);
  if (beacon_output == NULL)
  {
    printf("Could not allocate beacon output.\n");
    return -4;
  }

  if (libwifi_dump_beacon(&beacon, beacon_output, beacon_size) < 0)
  {
    printf("Could not dump beacon.\n");
    return -5;
  }

  libwifi_free_beacon(&beacon);

  *buf = beacon_output;

  return beacon_size;
}

int prepend_radiotap(const unsigned char *input, int input_size, unsigned char **output)
{
  static const unsigned char radiotap[] = "\x00\x00\x08\x00\x00\x00\x00\x00";
  static const int radiotap_size = 8;

  int output_size = input_size + radiotap_size;

  unsigned char *buf = (unsigned char *)malloc(output_size);
  if (buf == NULL)
  {
    return -1;
  }

  memcpy(buf, radiotap, radiotap_size);
  memcpy(buf + radiotap_size, input, input_size);

  *output = buf;

  return output_size;
}

int send_packet(int sock, int device_index, const unsigned char dst_addr[6], const unsigned char *packet, int packet_size)
{
  struct sockaddr_ll socket_address;
  socket_address.sll_ifindex = device_index;
  socket_address.sll_halen = ETH_ALEN;
  memcpy(socket_address.sll_addr, dst_addr, 6);

  pthread_mutex_lock(&socket_mutex);

  int ret = 0;
  if (sendto(sock, packet, packet_size, 0, (struct sockaddr *)&socket_address, sizeof(struct sockaddr_ll)) < 0)
  {
    printf("Could not send packet: %d\n", errno);
    ret = errno;
  }

  pthread_mutex_unlock(&socket_mutex);

  return ret;
}

int send_packet_with_radiotap(int sock, int device_index, const unsigned char dst_addr[6], const unsigned char *packet, int packet_size)
{
  struct libwifi_radiotap_info radiotap_info = {0};
  radiotap_info.present = (1 << IEEE80211_RADIOTAP_FLAGS) | (1 << IEEE80211_RADIOTAP_RATE);
  radiotap_info.flags = IEEE80211_RADIOTAP_F_FCS | IEEE80211_RADIOTAP_F_SHORTPRE;
  radiotap_info.rate_raw = 4;

  unsigned char radiotap_buffer[256];
  int radiotap_size = libwifi_create_radiotap(&radiotap_info, radiotap_buffer);
  
  uint32_t fcs = libwifi_calculate_fcs(packet, packet_size);

  int buffer_size = radiotap_size + packet_size + sizeof(uint32_t);
  unsigned char *buffer = malloc(buffer_size);
  if (buffer == NULL)
  {
    printf("Could not prepend radiotap.\n");
    return -2;
  }

  memcpy(buffer, radiotap_buffer, radiotap_size);
  memcpy(buffer + radiotap_size, packet, packet_size);
  memcpy(buffer + radiotap_size + packet_size, (unsigned char*)&fcs, sizeof(uint32_t));

  int ret = send_packet(sock, device_index, dst_addr, buffer, buffer_size);
  free(buffer);

  return ret;
}

int send_ack(int sock, int device_index, const unsigned char dst_addr[6])
{
  struct libwifi_cts ack_frame = {0};
  if (libwifi_create_cts(&ack_frame, dst_addr, 0)) {
    printf("Could not create ack packet.\n");
  }

  ack_frame.frame_header.frame_control.subtype = SUBTYPE_ACK;

  int ret = send_packet_with_radiotap(sock, device_index, dst_addr, (unsigned char*)&ack_frame, sizeof(struct libwifi_cts));

  return ret;
}

// output buffer size must be at least 18 bytes
void format_mac_address(const unsigned char addr[6], char *output)
{
  sprintf(output, "%02X:%02X:%02X:%02X:%02X:%02X", addr[0], addr[1], addr[2], addr[3], addr[4], addr[5]);
}

struct thread_info
{
  pthread_t thread_id;
  int sock;
  unsigned char tx_addr[6];
  int device_index;

  struct addr_list authenticated;
};

void handle_incoming_packet(u_char *userdata, const struct pcap_pkthdr *h, const u_char *bytes)
{
  struct thread_info *tinfo = (struct thread_info *)userdata;

  struct libwifi_frame frame = {0};
  if (libwifi_get_wifi_frame(&frame, bytes, h->caplen, true))
  {
    printf("Could not parse frame.\n");
    return;
  }

  if (frame.frame_control.type == TYPE_MANAGEMENT && frame.frame_control.subtype == SUBTYPE_AUTH)
  {
    if (strncmp(frame.header.data.addr1, tinfo->tx_addr, 6)) {
      // This is not for me. Misdelivery.
      return;
    }

    unsigned char in_addr[18];
    format_mac_address(frame.header.data.addr2, in_addr);
    printf("Authentication request from %s\n", in_addr);

    if (addr_list_contains(&tinfo->authenticated, frame.header.data.addr2))
    {
      // Already authenticated.
      return;
    }

    struct libwifi_auth_fixed_parameters *afp = (struct libwifi_auth_fixed_parameters *)frame.body;
    if (afp->algorithm_number != AUTH_OPEN || afp->status_code != STATUS_SUCCESS || afp->transaction_sequence != 1)
    {
      // Incorrect auth details.
      return;
    }

    //usleep(frame.header.ctrl.duration / 2 * 100);

    send_ack(tinfo->sock, tinfo->device_index, frame.header.data.addr2);

    addr_list_add(&tinfo->authenticated, frame.header.data.addr2);

    struct libwifi_auth auth_resp = {0};
    if (libwifi_create_auth(&auth_resp, frame.header.data.addr2, frame.header.data.addr1, frame.header.data.addr1, AUTH_OPEN, 2, STATUS_SUCCESS))
    {
      printf("Could not create auth response.\n");
      return;
    }

    auth_resp.frame_header.duration = frame.header.ctrl.duration;
    auth_resp.frame_header.seq_control.sequence_number = get_next_seq_number();

    //static const unsigned char broadcom_tag_value[] = {0x00, 0x10, 0x18, 0x02, 0x00, 0xf0, 0x00, 0x00, 0x00};
    //libwifi_quick_add_tag(&auth_resp.tags, TAG_VENDOR_SPECIFIC, broadcom_tag_value, 9);

    size_t auth_resp_size = libwifi_get_auth_length(&auth_resp);
    unsigned char *auth_resp_output = (unsigned char *)malloc(auth_resp_size);
    if (auth_resp_output == NULL)
    {
      printf("Could not allocate auth response output.\n");
      return;
    }

    if (libwifi_dump_auth(&auth_resp, auth_resp_output, auth_resp_size) < 0)
    {
      printf("Could not dump auth response.\n");
      return;
    }

    libwifi_free_auth(&auth_resp);

    //usleep(frame.header.ctrl.duration * 100);

    if (send_packet_with_radiotap(tinfo->sock, tinfo->device_index, frame.header.data.addr2, auth_resp_output, auth_resp_size))
    {
      printf("Could not send auth response.\n");
      return;
    }

    usleep(5000);

    if (send_packet_with_radiotap(tinfo->sock, tinfo->device_index, frame.header.data.addr2, auth_resp_output, auth_resp_size))
    {
      printf("Could not send auth response.\n");
      return;
    }

    printf("Successfully authenticated %s!\n", in_addr);
  }
}

void *polling_thread(void *arg)
{
  struct thread_info *tinfo = (struct thread_info *)arg;

  /*for (;;) {
    struct pollfd pfd;
    pfd.fd = tinfo->sock;
    pfd.events = POLLIN;

    if (poll(&pfd, 1, -1) < 0) {
      printf("Error polling.\n");
    } else {
      if (pfd.revents & POLLIN) {
        unsigned char message[1024];
        int msglen = read(tinfo->sock, message, 1024);
        if (msglen < 0) {
          printf("Could not read from socket.\n");
        } else {
          struct libwifi_frame frame = {0};
          if (libwifi_get_wifi_frame(&frame, message, msglen, true)) {
            printf("Could not parse frame.\n");
          } else {
            printf("Frame type: %d:%d\n", frame.frame_control.type, frame.frame_control.subtype);
          }
        }
      }
    }
  }*/

  char errbuf[PCAP_ERRBUF_SIZE];
  pcap_t *sniffer = pcap_create("mon0", errbuf);
  if (sniffer == NULL)
  {
    printf("Could not open sniffer: %s\n", errbuf);
    return NULL;
  }

  pcap_set_promisc(sniffer, 1);
  pcap_set_immediate_mode(sniffer, 1);
  pcap_activate(sniffer);
  pcap_loop(sniffer, -1, handle_incoming_packet, (u_char *)tinfo);

  return NULL;
}

int main(int argc, char **argv)
{
  static const char *interface_name = "mon0";

  int sock = socket(AF_PACKET, SOCK_RAW, 0);
  if (sock == -1)
  {
    printf("Could not open socket.\n");
    return 1;
  }

  unsigned char tx_addr[6];
  if (get_interface_mac_address(sock, interface_name, tx_addr))
  {
    printf("Could not get hardware address.\n");
    return 2;
  }

  int device_index = 0;
  if (get_interface_index(sock, interface_name, &device_index))
  {
    printf("Could not get device index.\n");
    return 2;
  }

  struct thread_info *tinfo = (struct thread_info *)malloc(sizeof(struct thread_info *));
  tinfo->sock = sock;
  memcpy(tinfo->tx_addr, tx_addr, 6);
  tinfo->device_index = device_index;
  addr_list_init(&tinfo->authenticated);

  pthread_attr_t thread_attr;
  if (pthread_attr_init(&thread_attr))
  {
    printf("Could not initialize thread attr.\n");
    return 6;
  }

  pthread_create(&tinfo->thread_id, &thread_attr, &polling_thread, tinfo);
  pthread_attr_destroy(&thread_attr);

  int beacon_index = 0;
  for (;;)
  {
    unsigned char *beacon_output;
    int beacon_size = make_beacon_frame(beacon_index, tx_addr, get_next_seq_number(), &beacon_output);
    if (beacon_size < 0)
    {
      return 3;
    }

    if (send_packet_with_radiotap(sock, device_index, kBroadcastAddress, beacon_output, beacon_size))
    {
      return 4;
    }

    free(beacon_output);

    beacon_index = (beacon_index + 1) % 10;

    usleep(1024 * 200);
  }

  if (pthread_join(tinfo->thread_id, NULL))
  {
    printf("Could not join thread I guess.\n");
    return 6;
  }

  free(tinfo);

  return 0;
}