431 lines
14 KiB
Python
431 lines
14 KiB
Python
"""Resumable batched build of the deduped query dataset.
|
|
|
|
Why batches: the raw collection is ~908 GB but the volume plus a sort temp
|
|
cannot hold a global sort. Each batch is parsed, sorted+deduped on the scratch
|
|
disk, and split into the same bucket tree; a final per-bucket sort dedups
|
|
across batches. Peak disk = source + tree.
|
|
|
|
Resume: completed batches are listed in ``build_state.json``; within a batch,
|
|
each parsed input file is recorded in ``progress_<n>.txt``. Stop with Ctrl-C /
|
|
``systemctl stop`` and rerun the same command to continue. The parse phase
|
|
resumes per file; sort+split of an interrupted batch is redone (duplicates are
|
|
removed by the final dedup), so an interrupted split does not lose data.
|
|
|
|
Archives are extracted into the scratch disk for the duration of their batch and
|
|
removed once parsed; the source is never modified.
|
|
"""
|
|
|
|
import bz2
|
|
import gzip
|
|
import json
|
|
import lzma
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import tarfile
|
|
import time
|
|
import zipfile
|
|
from pathlib import Path
|
|
|
|
from collectionparser.api import API, find_files, sort_unique
|
|
from collectionparser.parser import parse_files
|
|
from collectionparser.splitter import split_file
|
|
|
|
DEFAULT_BATCH_BYTES = 120 * 1024**3
|
|
ARCHIVE_SUFFIXES = (
|
|
".rar",
|
|
".zip",
|
|
".7z",
|
|
".tar",
|
|
".tar.gz",
|
|
".tgz",
|
|
".tar.bz2",
|
|
".tbz2",
|
|
".tar.xz",
|
|
".txz",
|
|
".gz",
|
|
".bz2",
|
|
".xz",
|
|
)
|
|
DEVNULL = subprocess.DEVNULL
|
|
|
|
|
|
def _is_archive(path: Path) -> bool:
|
|
return path.name.lower().endswith(ARCHIVE_SUFFIXES)
|
|
|
|
|
|
def _which(*names: str) -> str | None:
|
|
"""Return the path of the first available command among ``names``."""
|
|
for name in names:
|
|
found = shutil.which(name)
|
|
if found:
|
|
return found
|
|
return None
|
|
|
|
|
|
def _safe_target(dest: Path, name: str) -> Path:
|
|
target = (dest / name).resolve()
|
|
if not str(target).startswith(str(dest.resolve()) + os.sep):
|
|
raise ValueError(f"unsafe archive member: {name}")
|
|
return target
|
|
|
|
|
|
def _extract_compressed(path: Path, dest: Path) -> None:
|
|
"""Try the available RAR/7z extractors in turn.
|
|
|
|
No single tool wins: `unrar` handles real RARs, `7z` handles the many
|
|
mislabeled archives in the wild. Try both, first rc in (0, 1) wins.
|
|
ponytail: 7z may follow archive paths outside dest; trusted local dumps only.
|
|
"""
|
|
seven = _which("7z", "7za")
|
|
unrar = _which("unrar")
|
|
candidates = (
|
|
[unrar, seven] if path.name.lower().endswith(".rar") else [seven, unrar]
|
|
)
|
|
last = None
|
|
for tool in (c for c in candidates if c):
|
|
shutil.rmtree(dest, ignore_errors=True)
|
|
dest.mkdir(parents=True, exist_ok=True)
|
|
if "unrar" in tool:
|
|
cmd = [tool, "x", "-o+", "-idq", str(path), str(dest) + "/"]
|
|
else:
|
|
cmd = [tool, "x", "-y", f"-o{dest}", str(path)]
|
|
last = subprocess.run(
|
|
cmd, check=False, stdout=DEVNULL, stderr=DEVNULL
|
|
).returncode
|
|
if last in (0, 1):
|
|
return
|
|
shutil.rmtree(dest, ignore_errors=True)
|
|
raise RuntimeError(f"extractors failed (last rc={last})")
|
|
|
|
|
|
def _extract_tar(path: Path, dest: Path) -> None:
|
|
with tarfile.open(path) as t:
|
|
try:
|
|
t.extractall(dest, filter="data") # py3.12+
|
|
except TypeError: # py3.11
|
|
t.extractall(dest)
|
|
|
|
|
|
def _extract_archive(path: Path, dest: Path) -> None:
|
|
name = path.name.lower()
|
|
dest.mkdir(parents=True, exist_ok=True)
|
|
if name.endswith(".zip"):
|
|
with zipfile.ZipFile(path) as z:
|
|
for member in z.namelist():
|
|
_safe_target(dest, member)
|
|
z.extractall(dest)
|
|
elif name.endswith((".rar", ".7z")):
|
|
_extract_compressed(path, dest)
|
|
elif name.endswith(
|
|
(".tar", ".tar.gz", ".tgz", ".tar.bz2", ".tbz2", ".tar.xz", ".txz")
|
|
):
|
|
_extract_tar(path, dest)
|
|
else: # single-file .gz/.bz2/.xz
|
|
opener = (
|
|
gzip.open
|
|
if name.endswith(".gz")
|
|
else bz2.open
|
|
if name.endswith(".bz2")
|
|
else lzma.open
|
|
)
|
|
with opener(path, "rb") as r, open(dest / path.stem, "wb") as w:
|
|
shutil.copyfileobj(r, w, length=1 << 20)
|
|
|
|
|
|
def _discover(source: str) -> list[Path]:
|
|
files = find_files(source)
|
|
files.sort(key=lambda p: p.stat().st_size) # small files first
|
|
print(
|
|
f"{len(files)} files, {sum(1 for p in files if _is_archive(p))} archives.",
|
|
flush=True,
|
|
)
|
|
return files
|
|
|
|
|
|
def _batches(files: list[Path], batch_bytes: int):
|
|
batch: list[Path] = []
|
|
size = 0
|
|
for f in files:
|
|
sz = f.stat().st_size
|
|
if batch and size + sz > batch_bytes:
|
|
yield batch
|
|
batch, size = [], 0
|
|
batch.append(f)
|
|
size += sz
|
|
if batch:
|
|
yield batch
|
|
|
|
|
|
def _count_lines(path: str | Path) -> int:
|
|
total = 0
|
|
with open(path, "rb") as r:
|
|
for chunk in iter(lambda: r.read(1 << 20), b""):
|
|
total += chunk.count(b"\n")
|
|
return total
|
|
|
|
|
|
def _totals(stats: list[dict]) -> dict:
|
|
return {
|
|
"input_bytes": sum(s["input_bytes"] for s in stats),
|
|
"files": sum(s["files"] for s in stats),
|
|
"pairs": sum(s["pairs"] for s in stats),
|
|
"rejects": sum(s["rejects"] for s in stats),
|
|
"seconds": sum(s["total_s"] for s in stats),
|
|
}
|
|
|
|
|
|
def _stats_markdown(state: dict) -> str:
|
|
stats = state["stats"]
|
|
totals = _totals(stats)
|
|
lines = [
|
|
"# collectionparser build stats",
|
|
"",
|
|
f"- input: {totals['input_bytes'] / 1e9:.1f} GB across {totals['files']:,} files",
|
|
f"- valid pairs: {totals['pairs']:,}",
|
|
f"- rejects: {totals['rejects']:,}",
|
|
f"- batches: {len(stats)}",
|
|
f"- build seconds: {totals['seconds']:,.0f}",
|
|
]
|
|
if totals["seconds"]:
|
|
lines.append(
|
|
f"- aggregate throughput: {totals['input_bytes'] / 1e6 / totals['seconds']:.1f} MB/s"
|
|
)
|
|
if state.get("hashes") is not None:
|
|
lines.append(f"- hashed passwords collected: {state['hashes']:,}")
|
|
if state.get("dedup_s") is not None:
|
|
lines.append(f"- dedup seconds: {state['dedup_s']:,.0f}")
|
|
if state.get("unique_records") is not None:
|
|
lines.append(f"- unique records: {state['unique_records']:,}")
|
|
return "\n".join(lines) + "\n"
|
|
|
|
|
|
def _extract_batch(
|
|
archives: list[Path], extract_dir: Path, index: int
|
|
) -> tuple[list[Path], list[str]]:
|
|
"""Extract archives; return ``(files, failures)``.
|
|
|
|
Failures are logged, never silently dropped, so they can be recovered later.
|
|
"""
|
|
if not archives:
|
|
return [], []
|
|
t0 = time.time()
|
|
failures: list[str] = []
|
|
for i, archive in enumerate(archives):
|
|
try:
|
|
_extract_archive(archive, extract_dir / f"{i:05d}_{archive.name}")
|
|
except Exception as exc: # one bad archive must not stop the build
|
|
failures.append(f"{archive}\t{exc}")
|
|
files = [p for p in extract_dir.glob("**/*") if p.is_file()]
|
|
print(
|
|
f"batch {index}: {len(archives)} archives -> {len(files)} files, "
|
|
f"{len(failures)} failed ({time.time() - t0:.0f}s)",
|
|
flush=True,
|
|
)
|
|
for failure in failures[:5]:
|
|
print(f"WARN: {failure}", flush=True)
|
|
return files, failures
|
|
|
|
|
|
def build(
|
|
source: str,
|
|
dataset: str,
|
|
scratch: str,
|
|
batch_bytes: int = DEFAULT_BATCH_BYTES,
|
|
dedup: bool = True,
|
|
extract: bool = True,
|
|
) -> None:
|
|
"""Build the deduped query dataset from ``source`` into ``dataset``.
|
|
|
|
``scratch`` (a fast disk) holds per-batch files, sort temp and build state.
|
|
Safe to stop and rerun: completed batches are skipped and unparsed files of
|
|
the current batch are re-queued.
|
|
"""
|
|
dataset = Path(dataset).expanduser()
|
|
scratch = Path(scratch).expanduser()
|
|
dataset.mkdir(parents=True, exist_ok=True)
|
|
scratch.mkdir(parents=True, exist_ok=True)
|
|
os.environ["TMPDIR"] = str(scratch) # `sort -u` spills here, not into the dataset
|
|
for stale in scratch.glob("extract_*"):
|
|
shutil.rmtree(stale, ignore_errors=True)
|
|
|
|
if extract and not _which("7z", "7za", "unrar"):
|
|
print(
|
|
"WARN: no 7z/7za/unrar found; .rar/.7z archives will be skipped "
|
|
"(.zip/.tar/.gz still work).",
|
|
flush=True,
|
|
)
|
|
|
|
state_path = scratch / "build_state.json"
|
|
state = (
|
|
json.loads(state_path.read_text())
|
|
if state_path.exists()
|
|
else {"done": [], "stats": []}
|
|
)
|
|
|
|
def save():
|
|
state_path.write_text(json.dumps(state, indent=2))
|
|
|
|
done = set(state["done"])
|
|
files = _discover(source)
|
|
batches = list(_batches(files, batch_bytes))
|
|
state.update(
|
|
{
|
|
"batches_total": len(batches),
|
|
"batch_files": [len(b) for b in batches],
|
|
"dataset": str(dataset),
|
|
"complete": False,
|
|
}
|
|
)
|
|
save()
|
|
print(
|
|
f"{len(files)} files, {len(batches)} batches, {len(done)} already done.",
|
|
flush=True,
|
|
)
|
|
|
|
failures = scratch / "failures.txt"
|
|
extract_failures = scratch / "extract_failures.txt"
|
|
hashes = scratch / "hashes.txt"
|
|
for index, batch in enumerate(batches):
|
|
if index in done:
|
|
continue
|
|
state["phase"] = f"batch {index}"
|
|
save()
|
|
success = scratch / f"success_{index:05d}.txt"
|
|
failure = scratch / f"failure_{index:05d}.txt"
|
|
extract_dir = scratch / f"extract_{index:05d}"
|
|
progress_path = scratch / f"progress_{index:05d}.txt"
|
|
input_bytes = sum(f.stat().st_size for f in batch)
|
|
|
|
if not success.exists():
|
|
progress_path.unlink(missing_ok=True)
|
|
parsed = set()
|
|
if success.exists() and progress_path.exists():
|
|
with open(progress_path, encoding="utf8", errors="ignore") as r:
|
|
parsed = {line.rstrip("\n") for line in r}
|
|
|
|
inputs = [f for f in batch if not _is_archive(f)]
|
|
if extract:
|
|
extracted, failures_ = _extract_batch(
|
|
[f for f in batch if _is_archive(f)], extract_dir, index
|
|
)
|
|
inputs.extend(extracted)
|
|
if failures_:
|
|
with open(extract_failures, "a", encoding="utf8") as w:
|
|
w.write("\n".join(failures_) + "\n")
|
|
pending = [f for f in inputs if str(f) not in parsed]
|
|
|
|
t0 = time.time()
|
|
if pending:
|
|
parse_files(
|
|
pending,
|
|
str(success),
|
|
str(failure),
|
|
progress_filename=str(progress_path),
|
|
hashes_filename=str(hashes),
|
|
)
|
|
if success.exists():
|
|
pairs = _count_lines(success)
|
|
rejects = _count_lines(failure) if failure.exists() else 0
|
|
sort_unique(success, success)
|
|
split_file(str(success), str(dataset))
|
|
success.unlink(missing_ok=True)
|
|
if failure.exists():
|
|
with open(failures, "ab") as w, open(failure, "rb") as r:
|
|
shutil.copyfileobj(r, w, length=1024 * 1024)
|
|
failure.unlink(missing_ok=True)
|
|
progress_path.unlink(missing_ok=True)
|
|
else:
|
|
pairs = rejects = 0
|
|
shutil.rmtree(extract_dir, ignore_errors=True)
|
|
|
|
elapsed = time.time() - t0
|
|
stat = {
|
|
"batch": index,
|
|
"files": len(batch),
|
|
"input_bytes": input_bytes,
|
|
"pairs": pairs,
|
|
"rejects": rejects,
|
|
"total_s": round(elapsed, 1),
|
|
"mb_per_s": round(input_bytes / 1e6 / elapsed, 2) if elapsed else 0.0,
|
|
}
|
|
state["stats"].append(stat)
|
|
done.add(index)
|
|
state["done"] = sorted(done)
|
|
save()
|
|
(scratch / "stats.md").write_text(_stats_markdown(state))
|
|
print("BATCH " + json.dumps(stat), flush=True)
|
|
|
|
if hashes.exists():
|
|
sort_unique(hashes, hashes) # de-duplicate collected hashed passwords
|
|
state["hashes"] = _count_lines(hashes)
|
|
save()
|
|
|
|
if dedup and not state.get("deduped"):
|
|
state["phase"] = "deduping"
|
|
save()
|
|
t0 = time.time()
|
|
API.sort(str(dataset)) # per-bucket sort -u dedups across batches
|
|
state["dedup_s"] = round(time.time() - t0, 1)
|
|
state["deduped"] = True
|
|
save()
|
|
|
|
if state.get("unique_records") is None:
|
|
state["unique_records"] = sum(
|
|
_count_lines(p) for p in dataset.glob("**/*") if p.is_file()
|
|
)
|
|
save()
|
|
|
|
state["phase"] = "complete"
|
|
state["complete"] = True
|
|
save()
|
|
(scratch / "stats.md").write_text(_stats_markdown(state))
|
|
print("BUILD COMPLETE", flush=True)
|
|
|
|
|
|
def build_status(scratch: str) -> dict:
|
|
"""Return a compact, machine-readable snapshot of a build under ``scratch``."""
|
|
scratch = Path(scratch).expanduser()
|
|
state_path = scratch / "build_state.json"
|
|
if not state_path.exists():
|
|
return {"error": f"no build_state.json in {scratch}"}
|
|
state = json.loads(state_path.read_text())
|
|
done = state.get("done", [])
|
|
stats = state.get("stats", [])
|
|
total = state.get("batches_total")
|
|
totals = _totals(stats)
|
|
result = {
|
|
"phase": state.get("phase", "unknown"),
|
|
"complete": state.get("complete", False),
|
|
"batches_done": len(done),
|
|
"batches_total": total,
|
|
"deduped": state.get("deduped", False),
|
|
"input_gb": round(totals["input_bytes"] / 1e9, 2),
|
|
"pairs": totals["pairs"],
|
|
"rejects": totals["rejects"],
|
|
"hashes": state.get("hashes"),
|
|
"unique_records": state.get("unique_records"),
|
|
}
|
|
current = next((i for i in range(total or 0) if i not in done), None)
|
|
if current is not None:
|
|
progress = scratch / f"progress_{current:05d}.txt"
|
|
result["current_batch"] = current
|
|
result["current_files_done"] = (
|
|
sum(1 for _ in progress.open()) if progress.exists() else 0
|
|
)
|
|
batch_files = state.get("batch_files") or []
|
|
if current < len(batch_files):
|
|
result["current_files_total"] = batch_files[current]
|
|
for prefix in ("success", "failure"):
|
|
result[f"{prefix}_bytes"] = sum(
|
|
p.stat().st_size for p in scratch.glob(f"{prefix}_*.txt")
|
|
)
|
|
hashes_file = scratch / "hashes.txt"
|
|
result["hashes_bytes"] = hashes_file.stat().st_size if hashes_file.exists() else 0
|
|
dataset = state.get("dataset")
|
|
if dataset and Path(dataset).exists():
|
|
files = [p for p in Path(dataset).glob("**/*") if p.is_file()]
|
|
result["dataset_files"] = len(files)
|
|
result["dataset_bytes"] = sum(p.stat().st_size for p in files)
|
|
return result
|