195 lines
8.4 KiB
Python
195 lines
8.4 KiB
Python
"""Implementations behind the CLI commands: parse, split, sort, merge, query."""
|
|
|
|
import json
|
|
import os
|
|
import random
|
|
import shutil
|
|
import subprocess
|
|
from pathlib import Path
|
|
|
|
from tqdm import tqdm
|
|
|
|
from collectionparser.parser import parse_to_files
|
|
from collectionparser.query import QueryAPI
|
|
from collectionparser.splitter import get_output_file, write_line_to_file
|
|
from collectionparser.utils import create_dir_for_file, filter_printable, iter_lines, shuffle
|
|
|
|
# ponytail: buckets under this size sort byte-exactly in RAM; larger buckets
|
|
# shell out to `sort -u`, which spills to disk.
|
|
_IN_MEMORY_SORT_LIMIT = 8 * 1024 * 1024
|
|
|
|
|
|
class API:
|
|
"""Command implementations shared by the CLI."""
|
|
|
|
@staticmethod
|
|
def evaluate(old: str, new: str) -> None:
|
|
"""Report the hit rate of ``new`` against a sample of ``old``."""
|
|
hit_count = total_count = 0
|
|
files = [p for p in Path(old).expanduser().glob('**/*') if p.is_file()]
|
|
for _ in range(1000):
|
|
# reservoir sample 10 lines so arbitrarily large files stay O(1) in RAM.
|
|
random_lines: list[str] = []
|
|
with open(random.choice(files), encoding='utf8', errors='ignore') as reader:
|
|
for count, line in enumerate(reader):
|
|
if count < 10:
|
|
random_lines.append(line.rstrip('\n'))
|
|
elif (j := random.randint(0, count)) < 10:
|
|
random_lines[j] = line.rstrip('\n')
|
|
for random_line in random_lines:
|
|
password_list = QueryAPI.query(new, random_line.partition(':')[0], 'email')
|
|
total_count += 1
|
|
hit = 'hit' if password_list else 'miss'
|
|
hit_count += bool(password_list)
|
|
print(f'OLD: {random_line}, NEW: {hit}, HIT_RATE = {hit_count / total_count:.3f}.')
|
|
|
|
@staticmethod
|
|
def merge(src: str, dest: str) -> None:
|
|
"""Merge the bucket tree ``src`` into ``dest`` (destructive)."""
|
|
src_root = Path(src).expanduser()
|
|
dest_root = Path(dest).expanduser()
|
|
dest_files = {str(f).replace(str(dest_root), ''): f for f in find_files(dest_root)}
|
|
for src_file in find_files(src_root):
|
|
src_index = str(src_file).replace(str(src_root), '')
|
|
if src_index in dest_files:
|
|
append_file(src_file, dest_files[src_index])
|
|
else:
|
|
dst = dest_root / src_index.lstrip('/')
|
|
create_dir_for_file(str(dst))
|
|
move_file(src_file, dst)
|
|
|
|
@staticmethod
|
|
def parse(path: str, success_file: str, failure_file: str,
|
|
hashes_file: str | None = None) -> None:
|
|
parse_to_files(path, success_file, failure_file, hashes_file)
|
|
|
|
@staticmethod
|
|
def sort(path: str) -> None:
|
|
"""Sort and de-duplicate every file under ``path`` in place."""
|
|
files = [p for p in Path(path).expanduser().glob('**/*')
|
|
if p.is_file() and p.suffix != '.sorted']
|
|
for file in tqdm(files, desc='Unique > Sorting > Replacing existing files'):
|
|
output_path = file.with_suffix('.sorted')
|
|
sort_unique(file, output_path)
|
|
shutil.move(output_path, file)
|
|
|
|
@staticmethod
|
|
def clean(path: str) -> None:
|
|
"""Move misplaced records to their correct bucket and sort each file."""
|
|
files = find_files(Path(path).expanduser())
|
|
shuffle(files)
|
|
num_files_moved = 0
|
|
handles: dict[str, object] = {}
|
|
with tqdm(files, desc='clean') as bar:
|
|
for file in bar:
|
|
key = str(file).replace(path, '').removeprefix('/')
|
|
unique_file = str(file) + '.unique'
|
|
# stream: moved lines go to their target file, kept lines to a
|
|
# temp file. Nothing is held in RAM.
|
|
with open(str(file), encoding='utf8', errors='ignore') as reader, \
|
|
open(unique_file, 'w', encoding='utf8') as writer:
|
|
for line in reader:
|
|
email = line.partition(':')[0]
|
|
if get_output_file(email) != key:
|
|
write_line_to_file(handles, line, Path(path) / get_output_file(email),
|
|
buffer_scale=10)
|
|
num_files_moved += 1
|
|
if num_files_moved % 100_000 == 0:
|
|
bar.set_description(f'clean. {num_files_moved:,} moved.')
|
|
else:
|
|
writer.write(line)
|
|
sort_unique(unique_file, file)
|
|
os.remove(unique_file)
|
|
|
|
@staticmethod
|
|
def test(file: str, dataset: str) -> None:
|
|
"""Look up every email in ``file`` and print the cracked rate."""
|
|
cracked = total = 0
|
|
for line in iter_lines(file):
|
|
results = QueryAPI.query(dataset, line.strip(), 'email')
|
|
if results:
|
|
passwords = [p.partition(':')[2] for p in results]
|
|
print('FOUND:', line, ':'.join(passwords))
|
|
cracked += 1
|
|
total += 1
|
|
print(f'Cracked rate: {cracked / total if total else 0.0}')
|
|
|
|
@staticmethod
|
|
def chunk_files(path: str, max_size: int) -> None:
|
|
"""Split files larger than ``max_size`` MB and delete the originals."""
|
|
files = find_files(Path(path).expanduser())
|
|
for file in tqdm(files, desc='chunk large files'):
|
|
if file.stat().st_size / 1e6 > max_size:
|
|
# -d (numeric suffixes) is not compatible with MacOS.
|
|
cmd = ['split', '-b', str(int(max_size * 1e6)), str(file), str(file) + '_']
|
|
print(f'File size: {file.stat().st_size / 1e6:.2f} Mb. ' + ' '.join(cmd))
|
|
subprocess.run(cmd, check=True)
|
|
file.unlink()
|
|
|
|
@staticmethod
|
|
def export_jsonl(dataset: str, out: str) -> None:
|
|
"""Export the dataset as JSON Lines (`{"email":..., "password":...}`)."""
|
|
create_dir_for_file(out)
|
|
total = 0
|
|
with open(out, 'w', encoding='utf8') as writer:
|
|
for path in tqdm(find_files(Path(dataset).expanduser()), desc='export jsonl'):
|
|
with open(path, encoding='utf8', errors='ignore') as reader:
|
|
for line in reader:
|
|
email, _, password = line.rstrip('\n').partition(':')
|
|
writer.write(json.dumps({'email': email, 'password': password},
|
|
ensure_ascii=False) + '\n')
|
|
total += 1
|
|
print(f'Exported {total:,} records to {out}.')
|
|
|
|
@staticmethod
|
|
def filter_lines(file: str, out: str, less_than: int, more_than: int) -> None:
|
|
"""Keep lines whose printable length is within ``[less_than, more_than]``."""
|
|
create_dir_for_file(out)
|
|
with open(out, 'w', encoding='utf8') as writer, \
|
|
open(file, encoding='utf8', errors='ignore') as reader:
|
|
for line in reader:
|
|
printable = filter_printable(line)
|
|
if less_than <= len(printable.strip()) <= more_than:
|
|
writer.write(printable)
|
|
|
|
|
|
def find_files(path) -> list[Path]:
|
|
"""Return every regular file under ``path`` (recursively)."""
|
|
root = Path(path).expanduser()
|
|
return [p for p in root.glob('**/*') if p.is_file()]
|
|
|
|
|
|
def sort_unique(src, dst) -> None:
|
|
"""Write the sorted, de-duplicated lines of ``src`` to ``dst``.
|
|
|
|
In-memory for small files, ``sort -u`` for large ones. ``src`` may equal
|
|
``dst`` (in-place).
|
|
"""
|
|
src = Path(src)
|
|
if src.stat().st_size <= _IN_MEMORY_SORT_LIMIT:
|
|
lines = src.read_bytes().split(b'\n')
|
|
if lines and lines[-1] == b'':
|
|
lines.pop()
|
|
lines = sorted(set(lines))
|
|
Path(dst).write_bytes(b'\n'.join(lines) + b'\n' if lines else b'')
|
|
else:
|
|
tmp = f'{dst}.tmp'
|
|
with open(tmp, 'wb') as writer:
|
|
subprocess.run(['sort', '-u', str(src)], stdout=writer, check=True,
|
|
env={**os.environ, 'LC_ALL': 'C'})
|
|
os.replace(tmp, dst)
|
|
|
|
|
|
def append_file(src, dst) -> None:
|
|
"""Append ``src`` to ``dst`` and delete ``src`` (streams, constant memory)."""
|
|
print(f'merge {src} {dst}')
|
|
with open(src, 'rb') as reader, open(dst, 'ab') as writer:
|
|
shutil.copyfileobj(reader, writer, length=1024 * 1024)
|
|
print(f'del {src}')
|
|
os.remove(src)
|
|
|
|
|
|
def move_file(src, dst) -> None:
|
|
print(f'mv {src} {dst}')
|
|
shutil.move(src, dst)
|