"""
Scramblr: A single-file utility for secure passwords and encryption.
Author: Ethan Wu
License: MIT
"""

import os
import secrets
import string
import hashlib
import base64
import argparse

try:
    from cryptography.hazmat.primitives.ciphers.aead import AESGCM
    HAS_CRYPTO = True
except ImportError:
    HAS_CRYPTO = False

class Crypto:
    def __init__(self, iterations: int = 600_000):
        self.iterations = iterations

    # --- PASSWORDS ---
    def generate_password(self, length: int = 32) -> str:
        if length < 12: length = 12
        sets = [string.ascii_lowercase, string.ascii_uppercase, string.digits, "!@#$%^&*()"]
        pw = [secrets.choice(s) for s in sets]
        pool = "".join(sets)
        pw += [secrets.choice(pool) for _ in range(length - len(pw))]
        secrets.SystemRandom().shuffle(pw)
        return "".join(pw)

    def hash_password(self, password: str) -> str:
        salt = os.urandom(16)
        h = hashlib.pbkdf2_hmac("sha256", password.encode(), salt, self.iterations)
        return f"{salt.hex()}${h.hex()}"

    def verify_password(self, password: str, stored_hash: str) -> bool:
        salt_hex, hash_hex = stored_hash.split("$")
        attempt = hashlib.pbkdf2_hmac("sha256", password.encode(), bytes.fromhex(salt_hex), self.iterations)
        return secrets.compare_digest(attempt, bytes.fromhex(hash_hex))

    # --- AES-GCM (PRO) ---
    def encrypt(self, data: str, key_hex: str) -> str:
        if not HAS_CRYPTO: raise ImportError("Run 'pip install cryptography' for AES features.")
        aes = AESGCM(bytes.fromhex(key_hex))
        nonce = os.urandom(12)
        blob = aes.encrypt(nonce, data.encode(), None)
        return base64.b64encode(nonce + blob).decode()

    def decrypt(self, b64_data: str, key_hex: str) -> str:
        if not HAS_CRYPTO: raise ImportError("Run 'pip install cryptography' for AES features.")
        raw = base64.b64decode(b64_data)
        aes = AESGCM(bytes.fromhex(key_hex))
        return aes.decrypt(raw[:12], raw[12:], None).decode()

def main():
    parser = argparse.ArgumentParser(description="CryptoLab CLI")
    parser.add_argument("--gen-pw", type=int, help="Generate a password of specified length")
    parser.add_argument("--hash", type=str, help="Hash a password")
    parser.add_argument("--gen-key", action="store_true", help="Generate a random 256-bit AES key")
    
    args = parser.parse_args()
    lib = Crypto()

    if args.gen_pw:
        print(f"Generated Password: {lib.generate_password(args.gen_pw)}")
    elif args.hash:
        print(f"Hashed: {lib.hash_password(args.hash)}")
    elif args.gen_key:
        print(f"AES Key (Hex): {secrets.token_hex(32)}")
    else:
        parser.print_help()

if __name__ == "__main__":
    main()