collectionparser/collectionparser/cli.py

"""Command line entry point."""

import json

import click

from collectionparser.api import API
from collectionparser.parser import parse_to_files
from collectionparser.query import QueryAPI
from collectionparser.splitter import split_file
from collectionparser.utils import input_dir, input_file, output_dir, output_file


@click.group()
def cli():
    pass


@cli.command(
    short_help="Converts a large FILE to a query friendly folder OUT (e.g. a/b/c). "
    "Use RESTART_FROM to resume from the i-th line."
)
@click.option("--file", required=True, type=input_file())
@click.option("--out", required=True, type=output_dir())
@click.option("--restart_from", default=0, type=int, show_default=True)
def split(file, out, restart_from):
    split_file(file, out, restart_from)


@cli.command(
    short_help="Sorts and de-duplicates every file in a query friendly folder PATH."
)
@click.option("--path", required=True, type=input_dir())
def sort(path):
    API.sort(path)


@cli.command(
    short_help="Parses an unstructured folder PATH into SUCCESS_FILE, FAILURE_FILE "
    "and (optionally) HASHES_FILE."
)
@click.option("--path", required=True, type=input_dir())
@click.option("--success_file", required=True, type=output_file())
@click.option("--failure_file", required=True, type=output_file())
@click.option(
    "--hashes_file",
    required=False,
    type=output_file(),
    help="Collect hashed passwords here for later `crack`.",
)
def parse(path, success_file, failure_file, hashes_file):
    parse_to_files(path, success_file, failure_file, hashes_file)


@cli.command(short_help="Merges dataset SRC into dataset DEST, then sorts/dedups DEST.")
@click.option("--src", required=True, type=input_dir())
@click.option("--dest", required=True, type=input_dir())
def merge(src: str, dest: str):
    API.merge(src, dest)
    API.sort(dest)


@cli.command(
    short_help="Searches a dataset by exact email or wildcard pattern "
    '(e.g. "*@gmail.com", "john*", "jo?n@*").'
)
@click.option("--dataset", required=True, type=input_dir())
@click.option("--pattern", required=True)
def query(dataset: str, pattern: str):
    if "*" in pattern or "?" in pattern:
        for line in QueryAPI.wildcard(dataset, pattern):
            click.echo(line)
    else:
        results = QueryAPI.query(dataset, pattern)
        if results:
            click.echo("\n".join(results))


@cli.command(short_help="Exports a dataset as JSON Lines for OpenSearch/bulk ingest.")
@click.option("--dataset", required=True, type=input_dir())
@click.option("--out", required=True, type=output_file())
def export(dataset: str, out: str):
    API.export_jsonl(dataset, out)


@cli.command(short_help="Dictionary-attacks collected MD5/SHA1/SHA256 hashes.")
@click.option("--hashes", "hashes_path", required=True, type=input_file())
@click.option(
    "--wordlist",
    multiple=True,
    type=click.Path(),
    help="Plain wordlist(s); missing files are downloaded from rockyou.",
)
@click.option(
    "--passwords-from",
    multiple=True,
    type=click.Path(),
    help="Dataset file/folder whose plaintext passwords join the wordlist.",
)
@click.option("--out", required=True, type=output_file())
@click.option("--workers", default=0, type=int, show_default=True, help="0 = all CPUs.")
def crack(
    hashes_path: str, wordlist: tuple, passwords_from: tuple, out: str, workers: int
):
    from collectionparser.crack import crack as run_crack

    run_crack(
        hashes_path,
        out,
        wordlists=wordlist,
        passwords_from=passwords_from,
        workers=workers or None,
    )


@cli.command(
    short_help="Print build/dataset progress as JSON (for scripts and agents)."
)
@click.option(
    "--scratch",
    required=True,
    type=click.Path(exists=True, file_okay=False, dir_okay=True),
)
@click.option("--pretty", is_flag=True, help="Indent the JSON.")
def status(scratch: str, pretty: bool):
    from collectionparser.build import build_status

    click.echo(
        json.dumps(build_status(scratch), indent=2 if pretty else None, default=str)
    )


@cli.command(
    short_help="Builds the deduped query dataset from a raw source folder in batches. "
    "Stop and rerun to resume."
)
@click.option("--source", required=True, type=input_dir())
@click.option("--dataset", required=True, type=output_dir())
@click.option("--scratch", required=True, type=output_dir())
@click.option("--batch_gb", default=120, type=int, show_default=True)
@click.option("--dedup/--no-dedup", default=True)
@click.option(
    "--extract/--no-extract",
    default=True,
    show_default=True,
    help="Extract archives (.rar/.zip/.7z/.tar.*/.gz) found in the source.",
)
def build(
    source: str, dataset: str, scratch: str, batch_gb: int, dedup: bool, extract: bool
):
    from collectionparser.build import build as run_build

    run_build(
        source,
        dataset,
        scratch,
        batch_bytes=batch_gb * 1024**3,
        dedup=dedup,
        extract=extract,
    )


if __name__ == "__main__":
    cli()