import time
import os
import shutil
import subprocess
from django.db import transaction

from IoTAnalyzer.db_create import create_analyzed_hash
from IoTAnalyzer.models import HashRaw, HashAnalyzed, Tool, HashType, AnalyzedDataset

WORDLIST_PATH = "/app/django_project/IoTAnalyzer/management/commands/rockyou.txt"
HASHCAT_PATH = "hashcat"  # Update this if hashcat is in a different location

def crack_hashes_from_db(raw_hash_id):
    """ Crack hashes from the database using hashcat """

    if not os.path.exists(WORDLIST_PATH):
        return {"error": f"The wordlist file {WORDLIST_PATH} does not exist."}

    if not shutil.which(HASHCAT_PATH):
        return {"error": "Hashcat is not installed or not in PATH."}

    results = {}

    try:
        raw_hash = HashRaw.objects.get(id=raw_hash_id)
        hash_value = raw_hash.hash_value.strip()

        # Identify hash type and mode
        mode_info = identify_hash_mode(hash_value)
        if mode_info is None:
            return {"error": f"Unknown hash type for hash: {hash_value}"}

        mode, hash_type = mode_info
        print(f"Processing hash: {hash_value} (Type: {hash_type}, Mode: {mode})")

        temp_file_path = f"temp_hash_{raw_hash_id}.txt"
        with open(temp_file_path, "w", encoding="utf-8") as temp_file:
            temp_file.write(hash_value + "\n")

        # Run hashcat
        crack_command = [
            HASHCAT_PATH, "-m", str(mode), "-a", "0", "-D", "1,2", temp_file_path, WORDLIST_PATH, "--potfile-path", "hashcat.potfile", "--force"
        ]
        
        start_time = time.time()
        process = subprocess.run(crack_command, text=True, capture_output=True)
        end_time = time.time()
        duration = end_time - start_time
        crack_time_str = f"{duration:.2f}s"

        print("Hashcat Output:", process.stdout)
        print("Hashcat Errors:", process.stderr)

        analyzed_dataset, _ = AnalyzedDataset.objects.get_or_create(file=raw_hash.raw_dataset.file)
        tool, _ = Tool.objects.get_or_create(tool_name="Hashcat", defaults={"version": "6.x.x"})

        # Retrieve cracked hashes
        show_command = [
            HASHCAT_PATH, "-m", str(mode), "--show", "--potfile-path", "hashcat.potfile", temp_file_path
        ]
        show_process = subprocess.run(show_command, text=True, capture_output=True)
        print("Hashcat Show Output:", show_process.stdout)
        print("Hashcat Show Errors:", show_process.stderr)

        with transaction.atomic():
            if show_process.returncode == 0 and show_process.stdout.strip():
                cracked_hashes = show_process.stdout.strip().split("\n")
                for line in cracked_hashes:
                    parts = line.split(":", 1)  # Split only on the first colon
                    if len(parts) == 2:
                        full_hash, plaintext = parts
                        results[full_hash] = {"plaintext": plaintext, "type": hash_type}
                        print(f"Cracked Hash: {full_hash} -> {plaintext}")

                        create_analyzed_hash(analyzed_dataset.id, tool.id, hash_type, full_hash, True, plaintext, crack_time=crack_time_str)

                        print(
                            f"DEBUG: Saved hash {full_hash}, Cracked Value: {plaintext}")

            else:
                create_analyzed_hash(analyzed_dataset.id, tool.id, hash_type, hash_value, False, crack_time=crack_time_str)
                print(f"DEBUG: Saved uncracked hash {hash_value} ")

        os.remove(temp_file_path)

    except HashRaw.DoesNotExist:
        return {"error": f"No HashRaw found with ID: {raw_hash_id}"}
    except FileNotFoundError:
        return {"error": "Hashcat is not installed or not in PATH."}
    except Exception as e:
        return {"error": f"Error running hashcat: {str(e)}"}

    return results

def identify_hash_mode(hash_value):
    """
    Identify hashcat mode based on the hash length and known structures.
    """
    hash_modes = {
        32: (0, "md5"),
        40: (100, "sha1"),
        56: (300, "sha224"),
        64: (1400, "sha256"),
        96: (6100, "sha384"),
        128: (1700, "sha512")
    }
    return hash_modes.get(len(hash_value), None)