345 lines
14 KiB
Python
345 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 _stats_markdown(state: dict) -> str:
|
|
stats = state['stats']
|
|
total_in = sum(s['input_bytes'] for s in stats)
|
|
total_pairs = sum(s['pairs'] for s in stats)
|
|
total_rej = sum(s['rejects'] for s in stats)
|
|
total_s = sum(s['total_s'] for s in stats)
|
|
lines = [
|
|
'# collectionparser build stats',
|
|
'',
|
|
f'- input: {total_in / 1e9:.1f} GB across {sum(s["files"] for s in stats):,} files',
|
|
f'- valid pairs: {total_pairs:,}',
|
|
f'- rejects: {total_rej:,}',
|
|
f'- batches: {len(stats)}',
|
|
f'- build seconds: {total_s:,.0f}',
|
|
]
|
|
if total_s:
|
|
lines.append(f'- aggregate throughput: {total_in / 1e6 / total_s:.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': []}
|
|
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})
|
|
state_path.write_text(json.dumps(state, indent=2))
|
|
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}'
|
|
state_path.write_text(json.dumps(state, indent=2))
|
|
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)
|
|
state_path.write_text(json.dumps(state, indent=2))
|
|
(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)
|
|
state_path.write_text(json.dumps(state, indent=2))
|
|
|
|
if dedup and not state.get('deduped'):
|
|
state['phase'] = 'deduping'
|
|
state_path.write_text(json.dumps(state, indent=2))
|
|
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
|
|
state_path.write_text(json.dumps(state, indent=2))
|
|
|
|
if state.get('unique_records') is None:
|
|
state['unique_records'] = sum(_count_lines(p) for p in dataset.glob('**/*') if p.is_file())
|
|
state_path.write_text(json.dumps(state, indent=2))
|
|
|
|
state['phase'] = 'complete'
|
|
state['complete'] = True
|
|
state_path.write_text(json.dumps(state, indent=2))
|
|
(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')
|
|
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(sum(s['input_bytes'] for s in stats) / 1e9, 2),
|
|
'pairs': sum(s['pairs'] for s in stats),
|
|
'rejects': sum(s['rejects'] for s in stats),
|
|
'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
|