import hashlib
import itertools
import time
import string

# Define character set: 0-9, a-z, A-Z, and special symbols
CHARACTERS = string.digits + string.ascii_lowercase + string.ascii_uppercase + '#$%^&*'


def generate_salted_sha256_hash(password, salt_bytes):
    """
    Generate SHA256 hash of a password concatenated with salt.
    
    Args:
        password: The password string
        salt_bytes: The salt as bytes
    
    Returns:
        SHA256 hash as hexadecimal string
    """
    # Concatenate password bytes and salt bytes, then hash
    combined = password.encode() + salt_bytes
    return hashlib.sha256(combined).hexdigest()


def read_salted_password_database(filename):
    """
    Read salted password database from a text file.
    
    Args:
        filename: Path to the file containing [username, salt, hash] entries
    
    Returns:
        List of [username, salt, hash] entries
    """
    database = []
    
    try:
        with open(filename, 'r') as f:
            for line in f:
                line = line.strip()
                if line and line.startswith('[') and line.endswith(']'):
                    # Remove brackets and split by comma
                    content = line[1:-1]
                    parts = content.split(',', 2)  # Split into max 3 parts
                    if len(parts) == 3:
                        username = parts[0].strip()
                        salt_hex = parts[1].strip()
                        password_hash = parts[2].strip()
                        # Convert salt from hex string to bytes
                        salt_bytes = bytes.fromhex(salt_hex)
                        database.append([username, salt_bytes, password_hash])
        
        print(f"Loaded {len(database)} accounts from {filename}")
        return database
    
    except FileNotFoundError:
        print(f"Error: File '{filename}' not found.")
        return []
    except Exception as e:
        print(f"Error reading file: {e}")
        return []


def dictionary_attack_salted(password_database, N):
    """
    Perform dictionary attack on salted password database.
    
    Args:
        password_database: List of [username, salt, hash] entries
        N: Length of passwords to try
    
    Returns:
        Dictionary mapping usernames to cracked passwords
        Time taken for the attack
    """
    print(f"\n{'='*60}")
    print(f"Starting Dictionary Attack on Salted Hashes")
    print(f"Password length: {N}")
    print(f"Character set size: {len(CHARACTERS)}")
    print(f"Total possible passwords: {len(CHARACTERS)**N}")
    print(f"Number of accounts to crack: {len(password_database)}")
    print(f"{'='*60}\n")
    
    cracked_passwords = {}
    start_time = time.time()
    attempts = 0
    
    # For each user, we need to try all possible passwords with their specific salt
    for username, salt_bytes, stored_hash in password_database:
        print(f"Attacking {username} (salt: {salt_bytes.hex()})...")
        user_cracked = False
        user_attempts = 0
        
        # Generate all possible passwords of length N
        for password_tuple in itertools.product(CHARACTERS, repeat=N):
            password = ''.join(password_tuple)
            user_attempts += 1
            attempts += 1
            
            # Hash the password with this user's salt
            password_hash = generate_salted_sha256_hash(password, salt_bytes)
            
            # Check if this hash matches the stored hash
            if password_hash == stored_hash:
                cracked_passwords[username] = password
                print(f"[+] Cracked! {username}: {password} (after {user_attempts} attempts)")
                user_cracked = True
                break
            
            # Progress indicator for this user
            if user_attempts % 10000 == 0:
                print(f"  Tried {user_attempts} passwords for {username}...")
        
        if not user_cracked:
            print(f"[-] Failed to crack {username}")
    
    end_time = time.time()
    time_taken = end_time - start_time
    
    print(f"\n{'='*60}")
    print(f"Attack Complete!")
    print(f"Time taken: {time_taken:.4f} seconds")
    print(f"Total attempts: {attempts}")
    print(f"Passwords cracked: {len(cracked_passwords)}/{len(password_database)}")
    print(f"{'='*60}\n")
    
    return cracked_passwords, time_taken


def main():
    """Main function to run the dictionary attack on salted passwords."""
    print("Dictionary Attack on Salted SHA256 Password Hashes")
    print("=" * 60)
    
    # Get input file and password length from user
    filename = input("Enter password database filename: ")
    N = int(input("Enter password length N: "))
    
    # Warn if N is too large
    if N > 4:
        total_passwords = len(CHARACTERS) ** N
        print(f"\nWarning: With N={N}, there are {total_passwords:,} possible passwords.")
        print("With salting, each user requires a separate full search!")
        confirm = input("Continue? (yes/no): ")
        if confirm.lower() != 'yes':
            print("Attack cancelled.")
            return
    
    # Read password database from file
    print("\nReading password database...")
    password_database = read_salted_password_database(filename)
    
    if not password_database:
        print("No valid data found. Exiting.")
        return
    
    print("\nPassword Database:")
    for username, salt_bytes, password_hash in password_database:
        print(f"[{username}, {salt_bytes.hex()}, {password_hash}]")
    
    # Perform dictionary attack
    cracked_passwords, time_taken = dictionary_attack_salted(password_database, N)
    
    # Display results
    print("\nRecovered Passwords:")
    for username, password in cracked_passwords.items():
        print(f"{username}: {password}")
    
    print(f"\nTotal time needed: {time_taken:.4f} seconds")


if __name__ == "__main__":
    main()