import gradio as gr
import spaces
import subprocess
import os

# --- SETUP ---
def get_hashcat_binary():
    binary_path = "./hashcat-7.1.2/hashcat.bin"
    
    if not os.path.exists("hashcat-7.1.2"):
        print("Installing dependencies...")
        subprocess.run(["apt-get", "update"], check=True)
        subprocess.run(["apt-get", "install", "-y", "p7zip-full", "wget"], check=True)
        
        print("Downloading hashcat 7.1.2...")
        subprocess.run(["wget", "https://hashcat.net/files/hashcat-7.1.2.7z"], check=True)
        subprocess.run(["7z", "x", "hashcat-7.1.2.7z"], check=True)
        subprocess.run(["chmod", "+x", binary_path], check=True)
    
    return binary_path

# --- CORE LOGIC ---
@spaces.GPU(duration=600) 
def run_hashcat(hash_input, wordlist_input, allowed_symbols, attack_mode):
    hc_bin = get_hashcat_binary()
    
    if os.path.exists("target.hash"): 
        os.remove("target.hash")
    if os.path.exists("custom.dict"): 
        os.remove("custom.dict")
    
    with open("target.hash", "w") as f:
        f.write(hash_input.strip() + "\n")
        
    with open("custom.dict", "w") as f:
        f.write(wordlist_input.strip())
    
    # Build command as list to avoid shell escaping issues
    if attack_mode == "Quick Check (NTLM)":
        cmd = [hc_bin, "-m", "1000", "-a", "0", "target.hash", "custom.dict", "--force", "--backend-ignore-cuda"]
    
    elif attack_mode == "Hybrid (Words + 4 Custom Chars)":
        cmd = [hc_bin, "-m", "1000", "-a", "6", "-1", allowed_symbols, "target.hash", "custom.dict", "?1?1?1?1", "--force", "--backend-ignore-cuda"]
    
    elif attack_mode == "Hybrid (Words + 6 Custom Chars)":
        cmd = [hc_bin, "-m", "1000", "-a", "6", "-1", allowed_symbols, "target.hash", "custom.dict", "?1?1?1?1?1?1", "--force", "--backend-ignore-cuda"]
    
    elif attack_mode == "Rules-Based Mutations":
        rules_dir = "/usr/share/hashcat/rules"
        rule_path = os.path.join(rules_dir, "best64.rule")
        
        if not os.path.exists(rule_path):
            return f"Error: Rule file not found at {rule_path}"
                
        cmd = [hc_bin, "-m", "1000", "-a", "0", "target.hash", "custom.dict", "-r", rule_path, "--force", "-O", "--backend-ignore-opencl"]
        
    else:
        return "Error: Invalid mode selected"
    
    print(f"Exec: {' '.join(cmd)}")
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=600)
        return result.stdout + "\n--- STDERR ---\n" + result.stderr
    except Exception as e:
        return f"Error: {str(e)}"

# --- INTERFACE ---
with gr.Blocks() as demo:
    gr.Markdown("# ⚡ DPAPI Password Recovery Tool")
    
    with gr.Row():
        txt_hash = gr.Textbox(label="1. NT Hash", value="3232d48683d6ffc3e7ff9e9052c9a0f3")
        txt_words = gr.Textbox(label="2. Password Wordlist (one per line)", lines=8, placeholder="Edward9945\nLove060516\nRed1red")
    
    with gr.Row():
        txt_symbols = gr.Textbox(
            label="3. Allowed Characters", 
            value="acdegilmnorstuvwxz0123456789!@#$&*.",
            info="Only these chars will be tested"
        )
        mode_select = gr.Dropdown(
            [
                "Quick Check (NTLM)", 
                "Hybrid (Words + 4 Custom Chars)", 
                "Hybrid (Words + 6 Custom Chars)",
                "Rules-Based Mutations"
            ], 
            label="4. Attack Mode", 
            value="Hybrid (Words + 6 Custom Chars)"
        )
    
    btn_run = gr.Button("🔓 Crack Password", variant="primary")
    output_log = gr.Code(label="Hashcat Output", language="shell")
    
    btn_run.click(
        fn=run_hashcat, 
        inputs=[txt_hash, txt_words, txt_symbols, mode_select], 
        outputs=output_log
    )

if __name__ == "__main__":
    demo.launch()