#!/usr/bin/env python3

# yo dawg i heard you like libraries
from __future__ import annotations
import random
import string
import yaml
from datetime import datetime
from pathlib import Path
from typing import List, Iterable, Set, Dict, Any
from flask import Flask, jsonify, request, Response
import json, time, hashlib

# System Variables
words = []
simple_words = []
HASH_FILE = "/opt/pwdgen/hash_record.txt"
password_hashes =set()
SPECIAL_SET = "!@#$%^&*(),.<>?~`;:|][}{=-+_"
WORDS_FILE = "dict.yaml"
password_types = [
    "generate_standard_password",
    "generate_windows_ad_password", 
    "generate_simple_password"
]

#################################################
# Hash Record Functions
#################################################
# Create the file if it doesn’t exist
if not HASH_FILE.exists():
    HASH_FILE.touch(exist_ok=True)

# Read the hashes
try:
    with HASH_FILE.open("r", encoding="utf-8") as fh:
        password_hashes = {line.strip() for line in fh if line.strip()}
except OSError as exc:
    # If we cannot read the file we fall back to an empty set; the
    # application will continue to work, but any persisted hashes
    # will be lost.
    print(f"Warning: cannot read hash file {HASH_FILE!s}: {exc}")
    password_hashes = set()

def save_hashes() -> None:
    # Write to a temporary file first
    tmp_path = HASH_FILE.with_name(HASH_FILE.name + ".tmp")
    try:
        with tmp_path.open("w", encoding="utf-8") as fh:
            for h in sorted(password_hashes):
                fh.write(f"{h}\n")
        # Replace the original file atomically
        tmp_path.replace(HASH_FILE)
    except OSError as exc:
        print(f"Error: could not write hash file {HASH_FILE!s}: {exc}")

#################################################
# YAML Handler
#################################################
with open(WORDS_FILE, "r", encoding="utf-8") as fh:
    data = yaml.safe_load(fh)

# Ensure the key exists and is a list; otherwise return empty list.
raw_words = data.get("words") if isinstance(data, dict) else None
if not isinstance(words, list):
    raise ValueError(f"YAML file {path!s} must contain a top-level 'words' list")

# Ensure the key exists and is a list; otherwise return empty list.
raw_simple_words = data.get("simple_words") if isinstance(data, dict) else None
if not isinstance(words, list):
    raise ValueError(f"YAML file {path!s} must contain a top-level 'words' list")

# Strip whitespace and keep only non-empty words
words = [w.strip() for w in raw_words if isinstance(w, str) and w.strip()]
simple_words = [w.strip() for w in raw_simple_words if isinstance(w, str) and w.strip()]
total_words = len(words) + len(simple_words)

# Declare Descriptor after loading words

password_function_descriptor= {
    "generate_standard_password": {
        "name": "Standard Password",
        "type": "0",
        "description": f"From word lists totalling {total_words} word this algorithm selects 60 of these words that are less than 13 characters. From these 60, it uses the one with the index of the current second, and depending on this word length, it may or may not select additional words from the 60. It then randomly selects an integer and some special characters and randomly puts these all together into one string that is this password."
    }
    ,
    "generate_windows_ad_password": {
        "name": "Windows AD Password",
        "type": "1",
        "description": "This password is always in the following format: $Word1Word2Number$<br> - Each word is less than 7 characters, the number is 3 digits, and the $ represents a Special Charater."
    }
    , 
    "generate_simple_password": {
        "name": "Simple Password",
        "type": "2",
        "description": f"This simple password is in the following format: !Password123 - this pulls from a list of {len(simple_words)} simple words."
    }
}

#################################################
# Flask Routes and Helpers
#################################################
app = Flask(__name__)

@app.route('/get_password', methods=['GET'])
def get_password():
    pwd_index = request.args.get('pwd_index')
    # Check for presence
    if pwd_index is None:
        print("Missing required query parameter 'pwd_index'")
        return jsonify({
            "error": "Missing required query parameter 'pwd_index'"
        }), 400  # 400 Bad Request
    # Check for int 
    try:
        index = int(pwd_index)
    except ValueError:
        print(f"Parameter 'index' must be an integer, got '{pwd_index}'")
        return jsonify({
            "error": f"Parameter 'index' must be an integer, got '{pwd_index}'"
        }), 400
    # Check bounds
    if index < 0 or index >= len(password_types):
        print(f"Index {index} is out of bounds (0 ≤ index < {len(password_types)})")
        return jsonify({
            "error": f"Index {index} is out of bounds (0 ≤ index < {len(password_types)})",
            "password_types": password_types
        }), 404  # 404 Not Found

    # Success - return the password
    return jsonify({
        "password": get_password_by_index(index)
    }), 200

@app.route('/verbose_password', methods=['GET'])
def verbose_password():
    pwd_index = request.args.get('pwd_index')
    # Check for presence
    if pwd_index is None:
        return jsonify({
            "error": "Missing required query parameter 'pwd_index'"
        }), 400  # 400 Bad Request
    # Check for int 
    try:
        index = int(pwd_index)
    except ValueError:
        return jsonify({
            "error": f"Parameter 'index' must be an integer, got '{pwd_index}'"
        }), 400
    # Check bounds
    if index < 0 or index >= len(password_types):
        return jsonify({
            "error": f"Index {index} is out of bounds (0 ≤ index < {len(password_types)})",
            "password_types": password_types
        }), 404  # 404 Not Found
    # Success - return the result
    result = {
        "password": get_password_by_index(index),
        "descriptor": password_function_descriptor[password_types[index]],
        "password count": len(password_hashes)
    }
    return jsonify(result), 200

@app.route('/custom_password', methods=['POST'])
def custom_password():
    if not request.is_json:
        return jsonify(error="Request body must be JSON"), 400

    payload = request.get_json()

    # Basic presence check - you can relax this if you want defaults
    required_keys = {"w_min", "w_max", "w_count", "s_char", "num_len"}
    missing = required_keys - payload.keys()
    if missing:
        return jsonify(error=f"Missing keys: {', '.join(missing)}"), 400

    try:
        password = generate_custom_password(payload)
    except ValueError as exc:
        return jsonify(error=str(exc)), 400

    return jsonify({
        "password": password
    }), 200 

@app.route('/get_types', methods=['GET'])
def get_types():
    return jsonify(password_function_descriptor)

@app.route('/get_count', methods=['GET'])
def get_count():
    password_count = len(password_hashes)
    result = {
        "total_passwords": password_count
    }
    return jsonify(result)
    

def get_password_by_index(pwd_index: int):
    if pwd_index < 0 or pwd_index >= len(password_types):
        raise IndexError(f"Index {pwd_index} is out of range for list of length {len(password_types)}")
    func_name = password_types[pwd_index]
    # Look up the actual function object in the module namespace.
    func = globals().get(func_name)
    if not callable(func):
        raise ValueError(f"Function name '{func_name}' does not refer to a callable")
    return func()

#################################################
# Password Generators
#################################################
# Standard Password Generator
# This is excessively complicated
def generate_standard_password() -> str:
    candidates = [w for w in words if len(w) < 13]
    if not candidates:
        raise ValueError("No words of length < 13 found")
    # Shuffle the list and grab the first 60
    random.shuffle(candidates)
    candidates = candidates[:60]
    # Grab a random-ish word based on the time
    sec = int(time.time())
    word1 = candidates[sec % len(candidates)]
    len1 = len(word1)
    # Start building the password
    password = ""
    # If it picked a short word, add more stuff
    if len1 < 5:
        word2 = random.choice(words)
        len2 = len(word2)
        # If this picked a really short word, add moar things
        if len2 < 4:
            word3 = random.choice(words)
            return type3(word1, word2, word3)
        password = type2(word1, word2)
    # If it picked a really long word, this will be fine
    if len1 > 8:
        password = type1(word1)
    # If the first word is medium-sized, get one more
    word2 = random.choice(words)
    password = type2(word1, word2)
    # Check, and return password
    if check_and_add_hash(password):
        return password
    else:
        return "Somehow a password was duplicated"

# Standard Password Helper Functions
# For one really long word
def type1(w1: str) -> str:
    num1 = rand_int(100, 9999)
    sym = rand_symbol_set(SPECIAL_SET, 3)
    components = [ucfirst(w1), str(num1)] + sym
    shuffle_list(components)
    return "".join(components)
# For two medium-sized words
def type2(w1: str, w2: str) -> str:
    num2 = rand_int(100, 9999)
    sym = rand_symbol_set(SPECIAL_SET, 3)
    components = [ucfirst(w1), ucfirst(w2), str(num2)] + sym
    shuffle_list(components)
    return "".join(components)
# At least two tiny words
def type3(w1: str, w2: str, w3: str) -> str:
    num3 = rand_int(100, 9999)
    sym = rand_symbol_set(SPECIAL_SET, 4)
    components = [ucfirst(w1), ucfirst(w2), ucfirst(w3), str(num3)] + sym
    shuffle_list(components)
    return "".join(components)    

# Windows AD Password
def generate_windows_ad_password() -> str:
    short_words = [w for w in simple_words if len(w) < 7]
    if len(short_words) < 2:
        raise ValueError("Need at least two words of length < 7")
    # Get two words
    w1, w2 = random.sample(short_words, 2)
    # Get two symbols and a 3-digit number
    sym_set = "!@#$%^&*()[]{}|-+<>?"
    symbols = rand_symbol_set(sym_set, 2)
    number = rand_int(100, 999)
    # Build, check, and return password
    password = f"{symbols[0]}{ucfirst(w1)}{ucfirst(w2)}{number}{symbols[1]}"
    if check_and_add_hash(password):
        return password
    else:
        return "Somehow a password was duplicated"

# Simple Password
def generate_simple_password() -> str:
    candidates = [w for w in words if 7 <= len(w) <= 12]
    if not candidates:
        raise ValueError("No words of length 7-12 found")
    # Get a random word, number, and SC
    word = random.choice(candidates)
    symbol = random.choice("!@#$%^&*()")
    number = rand_int(100, 999)
    # Build, check, and return password
    password = f"{symbol}{ucfirst(word)}{number}"
    if check_and_add_hash(password):
        return password
    else:
        return "Somehow a password was duplicated"


#################################################
# Custom Password Generator
#################################################

def generate_custom_password(params: Dict[str, Any]) -> str:

    # Helper function for range
    def _num_range(num_len: int) -> tuple[int, int] | None:
        mapping = {
            0: None,
            1: (0, 9),
            2: (10, 99),
            3: (100, 999),
            4: (1000, 9999),
            5: (10000, 99999),
            6: (100000, 999999),
            7: (1000000, 9999999),
            8: (10000000, 99999999),
        }
        return mapping.get(num_len)

    # 2.1  Pull and validate parameters
    w_min      = int(params.get('w_min', 3))
    w_max      = int(params.get('w_max', 10))
    w_count    = int(params.get('w_count', 2))
    s_char     = int(params.get('s_char', 2))
    num_len    = int(params.get('num_len', 3))
    words_raw  = list(params.get('words', []))

    # Basic sanity checks
    if not (3 <= w_min <= 10):
        raise ValueError("w_min must be between 3 and 10")
    if not (3 <= w_max <= 10):
        raise ValueError("w_max must be between 3 and 10")
    if w_min > w_max:
        raise ValueError("w_min cannot be greater than w_max")
    if not (1 <= w_count <= 5):
        raise ValueError("w_count must be between 1 and 5")
    if not (0 <= s_char <= 4):
        raise ValueError("s_char must be between 0 and 4")
    if not (0 <= num_len <= 8):
        raise ValueError("num_len must be between 0 and 8")

    # 2.2  Filter the word list by length
    candidates = [w for w in words if w_min <= len(w) <= w_max]
    if len(candidates) < w_count:
        raise ValueError(
            f"Not enough words of length {w_min}-{w_max}. "
            f"Need {w_count}, found {len(candidates)}."
        )

    # 2.3  Randomly pick words and capitalize first letter
    rng = random.SystemRandom()  # cryptographically secure PRNG
    selected_words = rng.sample(candidates, w_count)
    selected_words = [w.capitalize() for w in selected_words]

    # Generate the numeric block (if required)
    num_part = ""
    num_range = _num_range(num_len)
    if num_range:
        low, high = num_range
        num_part = str(rng.randint(low, high))

    # --------------------------------------------------------------------- #
    # Pick special characters
    # --------------------------------------------------------------------- #
    special_chars = "!@#$%^&*()[]{}|-+<>?"
    shuffled_specials = list(special_chars)
    rng.shuffle(shuffled_specials)
    specials = "".join(shuffled_specials[:s_char])

    # --------------------------------------------------------------------- #
    # Assemble all pieces into a list, shuffle, and join
    # --------------------------------------------------------------------- #
    all_parts: List[str] = selected_words
    if num_part:
        all_parts.append(num_part)
    if specials:
        # We insert specials as individual characters (mirroring PHP code)
        all_parts.extend(list(specials))

    rng.shuffle(all_parts)
    password = "".join(all_parts)

    return password

#################################################
# Password helper functions
#################################################

# Return the word with its first character capitalised
def ucfirst(word: str) -> str:
    return word[0].upper() + word[1:] if word else ""

def shuffle_list(items: List[str]) -> None:
    random.shuffle(items)

def rand_int(low: int, high: int) -> int:
    return random.randint(low, high)

def rand_symbol_set(symbols: str, n: int) -> List[str]:
    return random.sample(symbols, n)

# Return True when password is unused
def check_and_add_hash(password: str) -> bool:
    def sha256_hex(s: str) -> str:
        return hashlib.sha256(s.encode('utf-8')).hexdigest()
    hash_value = sha256_hex(password)
    if hash_value in password_hashes:
        return False
    # Hash was missing - add it and report that it was newly inserted.
    password_hashes.add(hash_value)
    save_hashes()
    return True

#################################################
# Main Function
#################################################

if __name__ == "__main__":

    ######################################
    # Flask API
    ######################################

    app.run(debug=False, host='0.0.0.0', port=5000)