from pwnagotchi import plugins
import logging
import subprocess
import os
import json
import sys
from multiprocessing.pool import ThreadPool

class QuickDic(plugins.Plugin):
    __author__ = 'silentree12th'
    __version__ = '1.5'
    __license__ = 'GPL3'
    __description__ = 'Run a quick dictionary scan against captured handshakes.'
    __dependencies__ = {
        'apt': ['aircrack-ng'],
    }
    __defaults__ = {
        'enabled': True,
        'wordlist_folder': '/home/pi/wordlists/',
        'progress_file': '/home/pi/quickdic_progress.json',
        'face': '(·ω·)',
    }

    def __init__(self):
        self.text_to_set = ""
        self.progress = self.load_progress()

    def on_loaded(self):
        logging.info('[quickdic] plugin loaded')
        self.options.setdefault('face', '(·ω·)')
        self.options.setdefault('wordlist_folder', '/home/pi/wordlists/')
        self.options.setdefault('progress_file', '/home/pi/quickdic_progress.json')
        self.options.setdefault('enabled', True)

        # Restart with PyPy for better performance
        if sys.executable.endswith("python3"):
            logging.info("[quickdic] Restarting with PyPy...")
            os.execv("/usr/bin/pypy3", ["pypy3"] + sys.argv)

        # Check if aircrack-ng is installed
        check = subprocess.run(["dpkg", "-l", "aircrack-ng"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        if b"aircrack-ng" in check.stdout:
            logging.info('[quickdic] aircrack-ng is installed')
        else:
            logging.warning('[quickdic] aircrack-ng is not installed!')

    def load_progress(self):
        """Load progress from the progress file."""
        try:
            with open(self.options['progress_file'], 'r') as f:
                return json.load(f)
        except FileNotFoundError:
            return {}

    def save_progress(self):
        """Save progress to the progress file."""
        with open(self.options['progress_file'], 'w') as f:
            json.dump(self.progress, f)

    def try_wordlist(self, args):
        """Attempt to crack the handshake using a specific wordlist."""
        wordlist, filename, bssid = args
        wl_path = os.path.join(self.options['wordlist_folder'], wordlist)
        output_file = f"{filename}.{wordlist}.cracked"

        cmd = ["aircrack-ng", filename, "-w", wl_path, "-l", output_file, "-q", "-b", bssid]

        try:
            proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=300)
            if proc.returncode == 0:
                with open(output_file, 'r') as f:
                    pwd = f.read().strip()
                return (wordlist, True, pwd)
            return (wordlist, False, None)
        except Exception as e:
            logging.error(f"[quickdic] Error processing {wordlist}: {str(e)}")
            return (wordlist, False, None)

    def on_handshake(self, agent, filename, access_point, client_station):
        """Handle a captured handshake."""
        display = agent.view()
        bssid = access_point['mac']

        # Skip if already processed
        if filename in self.progress:
            logging.info(f'[quickdic] Handshake {filename} already processed. Skipping...')
            return

        # Verify handshake
        result = subprocess.run(["aircrack-ng", filename], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        if b"1 handshake" not in result.stdout:
            logging.info('[quickdic] No valid handshake')
            return
        else:
            logging.info('[quickdic] Handshake confirmed')

        # Get wordlists
        wordlist_dir = self.options['wordlist_folder']
        try:
            files = sorted(os.listdir(wordlist_dir))
        except Exception as e:
            logging.error(f"[quickdic] Error reading wordlist folder: {e}")
            return

        wordlists = [f for f in files if f.endswith('.txt')]
        if not wordlists:
            logging.warning("[quickdic] No wordlist files found.")
            return

        # Process wordlists in parallel
        pool = ThreadPool(processes=2)
        args = [(f, filename, bssid) for f in wordlists]
        results = pool.map(self.try_wordlist, args)
        pool.close()
        pool.join()

        # Check results
        found = False
        for wordlist, success, pwd in results:
            if success:
                logging.info(f"[quickdic] Password found: {pwd}")
                self.text_to_set = f"Cracked password: {pwd}"
                display.set('face', self.options['face'])
                display.set('status', self.text_to_set)
                display.update(force=True)
                found = True
                break

        # Update progress
        self.progress[filename] = "found" if found else "not found"
        self.save_progress()

        if not found:
            logging.info("[quickdic] No password found in any wordlist.")

    def on_ui_update(self, ui):
        """Update the Pwnagotchi UI."""
        if self.text_to_set:
            ui.set('face', self.options['face'])
            ui.set('status', self.text_to_set)
            self.text_to_set = ""

    def on_unload(self, ui):
        """Handle plugin unload."""
        with ui._lock:
            logging.info('[quickdic] plugin unloaded')