"""Dictionary attacks over the unsalted hashes collected by ``parse``.

Candidate sources: plain wordlists, the password column of existing
``email:password`` datasets, and rockyou (downloaded on demand). They are merged
through ``sort -u`` so memory stays bounded for arbitrarily large inputs.

Only unsalted MD5/SHA1/SHA256 are attackable; bcrypt is intentionally not
supported (a wordlist attack on it is infeasible).

Memory: the target index is built once and shared with the workers by forking
(copy-on-write), so it costs one copy, not one per worker. It still scales with
the number of hashes, so pass ``max_hashes`` to stream the input in bounded
chunks on memory-limited hosts. ponytail: platforms without ``fork`` fall back
to rebuilding it in each worker (RAM x workers); a shared mmap index is the
upgrade path for enormous hash files.
"""

import gzip
import hashlib
import multiprocessing
import os
import shutil
import subprocess
import urllib.request
from concurrent.futures import ProcessPoolExecutor
from itertools import islice
from pathlib import Path

from collectionparser.api import find_files

DEFAULT_WORDLIST_URL = "https://weakpass.com/download/90/rockyou.txt.gz"

# hex digest length -> hashlib constructor
_ALGOS = {32: hashlib.md5, 40: hashlib.sha1, 64: hashlib.sha256}
_HEX_DIGITS = frozenset(b"0123456789abcdefABCDEF")

# Filled once in `crack`; forked workers inherit it instead of rebuilding it.
_TARGETS: dict[int, dict[bytes, list[bytes]]] | None = None


def download_rockyou(dest: Path) -> Path:
    """Download and gunzip the default rockyou wordlist if not already present."""
    dest = Path(dest)
    if dest.exists() and dest.stat().st_size:
        return dest
    dest.parent.mkdir(parents=True, exist_ok=True)
    print(f"Downloading wordlist {DEFAULT_WORDLIST_URL} ...", flush=True)
    with (
        urllib.request.urlopen(DEFAULT_WORDLIST_URL) as response,
        gzip.GzipFile(fileobj=response) as gz,
        open(dest, "wb") as out,
    ):
        shutil.copyfileobj(gz, out, length=1 << 20)
    return dest


def _write_passwords(source: Path, out) -> None:
    """Append the password column of an ``email:password`` file to ``out``."""
    with open(source, encoding="utf8", errors="ignore") as reader:
        for line in reader:
            out.write(line.partition(":")[2].encode("utf8"))


def prepare_candidates(wordlists, passwords_from, work_dir: Path) -> Path:
    """Return a single, de-duplicated wordlist built from all candidate sources."""
    work_dir = Path(work_dir)
    work_dir.mkdir(parents=True, exist_ok=True)
    wordlists = [Path(w) for w in wordlists] or [work_dir / "rockyou.txt"]
    wordlists = [w if w.exists() else download_rockyou(w) for w in wordlists]
    if len(wordlists) == 1 and not passwords_from:
        return wordlists[0]

    unique = work_dir / "candidates.unique.txt"
    env = {**os.environ, "LC_ALL": "C", "TMPDIR": str(work_dir)}
    with (
        open(unique, "wb") as sink_file,
        # Match api.sort_unique: cap the sort buffer so a huge --passwords-from
        # dataset spills to disk instead of OOMing the host.
        subprocess.Popen(
            ["sort", "-S", "1G", "--parallel", "4", "-u"],
            stdin=subprocess.PIPE,
            stdout=sink_file,
            env=env,
        ) as process,
    ):
        with process.stdin as sink:
            for wordlist in wordlists:
                with open(wordlist, "rb") as reader:
                    shutil.copyfileobj(reader, sink, length=1 << 20)
            for source in passwords_from:
                root = Path(source).expanduser()
                for path in [root] if root.is_file() else find_files(root):
                    _write_passwords(path, sink)
        if process.wait() != 0:
            raise RuntimeError("failed to build the candidate wordlist")
    return unique


def load_targets(hashes_path: str | Path) -> dict[int, dict[bytes, list[bytes]]]:
    """Return ``{hex_length: {digest_bytes: [email, ...]}}`` from a hashes file."""
    targets: dict[int, dict[bytes, list[bytes]]] = {}
    with open(hashes_path, "rb") as reader:
        for raw in reader:
            line = raw.rstrip(b"\r\n")
            email, separator, digest = line.rpartition(b":")
            if (
                not separator
                or len(digest) not in _ALGOS
                or not _HEX_DIGITS.issuperset(digest)
            ):
                continue
            try:
                key = bytes.fromhex(digest.decode("ascii"))
            except ValueError:
                continue
            targets.setdefault(len(digest), {}).setdefault(key, []).append(email)
    return targets


def _scan_stripe(args: tuple) -> list[bytes]:
    """Attack one residue class of the wordlist (``index`` of every ``total``)."""
    hashes_path, wordlist_path, index, total = args
    targets = _TARGETS if _TARGETS is not None else load_targets(hashes_path)
    checks = [(targets[length], _ALGOS[length]) for length in targets]
    matches: list[bytes] = []
    with open(wordlist_path, "rb") as reader:
        for n, raw in enumerate(reader):
            if n % total != index:
                continue
            word = raw.rstrip(b"\r\n")
            if not word:
                continue
            for target, constructor in checks:
                emails = target.get(constructor(word).digest())
                if emails:
                    matches.extend(email + b":" + word + b"\n" for email in emails)
    return matches


def _iter_hash_chunks(hashes_path, max_hashes, work_dir):
    """Yield hash-file paths of at most ``max_hashes`` lines, or the whole file."""
    if not max_hashes or max_hashes <= 0:
        yield str(hashes_path)
        return
    index = 0
    with open(hashes_path, "rb") as reader:
        while True:
            lines = list(islice(reader, max_hashes))
            if not lines:
                break
            chunk = Path(work_dir) / f"hashes.chunk.{index:05d}.txt"
            chunk.write_bytes(b"".join(lines))
            try:
                yield str(chunk)
            finally:
                chunk.unlink(missing_ok=True)
            index += 1


def _attack(hashes_path, writer, wordlist, workers) -> int:
    """Attack one hashes file, writing ``email:password`` matches to ``writer``."""
    global _TARGETS
    _TARGETS = load_targets(hashes_path)
    targets = _TARGETS
    if not targets:
        print(f"No crackable MD5/SHA1/SHA256 hashes found in {hashes_path}.")
        return 0
    n_hashes = sum(
        len(emails) for by_digest in targets.values() for emails in by_digest.values()
    )
    print(
        f"Attacking {n_hashes:,} hashes with {wordlist.name} on {workers} worker(s).",
        flush=True,
    )

    jobs = [(str(hashes_path), str(wordlist), i, workers) for i in range(workers)]
    total = 0
    if workers == 1:
        results = map(_scan_stripe, jobs)
    else:
        # fork shares the parent's _TARGETS by copy-on-write.
        context = multiprocessing.get_context("fork") if hasattr(os, "fork") else None
        executor = ProcessPoolExecutor(max_workers=workers, mp_context=context)
        results = executor.map(_scan_stripe, jobs)
    try:
        for matches in results:
            for line in matches:
                writer.write(line)
                total += 1
    finally:
        if workers != 1:
            executor.shutdown()
    return total


def crack(
    hashes_path: str,
    out_path: str,
    wordlists=(),
    passwords_from=(),
    workers: int | None = None,
    work_dir: str | None = None,
    max_hashes: int | None = None,
) -> int:
    """Attack the hashes with the candidate wordlists; write ``email:password``.

    ``max_hashes`` bounds how many hash lines are held in the in-RAM target
    index at once; larger inputs are streamed in chunks that spill to ``work_dir``.
    A single global index of hundreds of millions of hashes needs tens of GB, so
    set this on memory-limited hosts. ponytail: the wordlist is rescanned per
    chunk, so a smaller ``max_hashes`` trades RAM for repeated passes.
    """
    work_dir = Path(work_dir) if work_dir else Path(out_path).parent
    wordlist = prepare_candidates(list(wordlists), list(passwords_from), work_dir)
    workers = max(1, workers or os.cpu_count() or 1)

    total = 0
    with open(out_path, "wb") as writer:
        for chunk_path in _iter_hash_chunks(hashes_path, max_hashes, work_dir):
            total += _attack(chunk_path, writer, wordlist, workers)
    print(f"Cracked {total:,} hashes -> {out_path}.")
    return total