#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 <stdio.h>
#include <sys/ioctl.h>
#include <unistd.h>

#include "beacon_data.h"

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

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 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, 7))
  {
    printf("Could not create beacon frame.\n");
    return -3;
  }

  beacon.frame_header.seq_control.sequence_number = seq_number;

  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 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);

  if (sendto(sock, packet, packet_size, 0, (struct sockaddr *)&socket_address, sizeof(struct sockaddr_ll)) < 0)
  {
    printf("Could not send packet.\n");
    return errno;
  }

  return 0;
}

int send_packet_with_radiotap(int sock, int device_index, const unsigned char dst_addr[6], const unsigned char *packet, int packet_size)
{
  unsigned char *buffer;
  int buffer_size = prepend_radiotap(packet, packet_size, &buffer);
  if (buffer_size < 0)
  {
    printf("Could not prepend radiotap.\n");
    return -2;
  }

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

  return ret;
}

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;
  }

  int beacon_index = 0;
  int seq_number = 0;
  for (;;)
  {
    unsigned char *beacon_output;
    int beacon_size = make_beacon_frame(beacon_index, tx_addr, 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;
    seq_number++;

    usleep(1024 * 100);
  }

  return 0;
}