"""Command-line interface for RandQuik.""" import argparse import gc import os import sys import time import aeg from randquik.benchmark import run_benchmark from randquik.crypto import derive_key, generate_random_seed from randquik.io import open_fd from randquik.progress import ProgressDisplay from randquik.utils import ( parse_size, print_summary, ) from randquik.workers import ( BLOCK_SIZE, FdProducer, ) __all__ = ["main"] # Disable GC for performance gc.disable() ciph: aeg.Cipher = None # type: ignore (set in main after parsing args) DEFAULT_ALG = "AEGIS-128X2" def prepare_seed(args): """Prepare seed and determine if it was generated.""" generated_seed = args.seed is None seed = generate_random_seed() if generated_seed else args.seed return seed, generated_seed def prepare_key(seed, keybytes): """Derive key from seed.""" key = derive_key(seed, keybytes) return key def parse_seeks(args): """Parse seek values from arguments.""" try: iseek = parse_size(args.iseek, args.output) or 0 oseek = parse_size(args.oseek, args.output) or 0 if args.seek: iseek = oseek = parse_size(args.seek, args.output) or 0 except ValueError as e: raise ValueError(f"Error parsing seek: {e}") from None return iseek, oseek def singlethreaded(args, total_bytes, oseek, start_time, key, seed_for_display): if args.verbose: mode_desc = "infinite output" if total_bytes is None else "workers=0" print( f"Special mode: {mode_desc} — single-buffer, single-threaded generation", file=sys.stderr, ) with open_fd( args.output, 0 if total_bytes is None else total_bytes, dry=args.dry, oseek=oseek ) as fd: written = 0 buf = bytearray(BLOCK_SIZE) view = memoryview(buf) nonce = bytearray(ciph.NONCEBYTES) progress_state = {"written": 0} progress = ProgressDisplay( total_bytes, start_time, progress_state, infinite=total_bytes is None, seed=seed_for_display, ) if not args.quiet: progress.start() try: while total_bytes is None or written < total_bytes: start = written if total_bytes is None: size = BLOCK_SIZE else: end = min(start + BLOCK_SIZE, total_bytes) size = end - start chunk = view[:size] ciph.stream(key, nonce, size, into=chunk) if not args.dry: os.write(fd, chunk) ciph.nonce_increment(nonce) written += size progress_state["written"] = written finally: progress.stop() elapsed = time.perf_counter() - start_time action = "generated" if args.dry else "wrote" if not args.quiet and total_bytes is not None: print_summary(written, elapsed, action, seed=seed_for_display) def _main(): """Internal main function that may raise exceptions.""" parser = argparse.ArgumentParser(description="Generate random bytes using AEGIS ciphers") parser.add_argument("-s", "--seed", help="Alphanumeric seed string", type=str) parser.add_argument( "-l", "--len", help="Length to generate (e.g. 1g, 100mi, 1000sect)", type=str, default=None, ) parser.add_argument("-o", "--output", help="Output file (default: stdout)", type=str) parser.add_argument( "-t", "--threads", help="Number of worker threads (benchmark: upper limit)", type=int, default=None, ) parser.add_argument( "-a", "--alg", help=f"Cipher algorithm (default: {DEFAULT_ALG})", type=str, default=None, ) parser.add_argument( "--benchmark", action="store_true", help="Run benchmark (generates 1GB and reports speed)", ) parser.add_argument( "--dry", action="store_true", help="Dry run: open output but skip writes (for benchmarking)", ) parser.add_argument( "--seek", type=str, default=None, help="Seek both input stream and output to position (e.g. 1g, 100mi)", ) parser.add_argument( "--iseek", type=str, default=None, help="Seek input random stream to position (overrides --seek)", ) parser.add_argument( "--oseek", type=str, default=None, help="Seek output file to position (overrides --seek)", ) parser.add_argument( "-q", "--quiet", action="store_true", help="Quiet mode: suppress all output except errors", ) parser.add_argument( "-v", "--verbose", action="store_true", help="Verbose mode: show I/O mode and timing statistics", ) args = parser.parse_args() # Normalize "-" output to None (stdout) if args.output == "-": args.output = None global ciph ciph = aeg.cipher(args.alg or DEFAULT_ALG) # Validate and process args seed, generated_seed = prepare_seed(args) key = prepare_key(seed, ciph.KEYBYTES) iseek, oseek = parse_seeks(args) total_bytes = parse_size(args.len) # None if not specified, 0 if -l0 # Seed hint for generated seeds seed_for_display = seed if generated_seed else None start_time = time.perf_counter() if args.benchmark: if args.seed is not None: raise ValueError("Cannot specify seed in benchmark mode") if iseek or oseek: raise ValueError("Cannot use seek options in benchmark mode") run_benchmark(args) return # Single-threaded mode (workers == 0) if args.threads == 0: return singlethreaded(args, total_bytes, oseek, start_time, key, seed_for_display) workers = args.threads if args.threads is not None else 1 # File mode with ring buffers with open_fd(args.output, total_bytes, dry=args.dry, oseek=oseek) as fd: producer = FdProducer( workers, key, ciph, total_bytes, fd, dry=args.dry, iseek=iseek, profile=args.verbose ) if args.verbose: dry_str = " (dry run)" if args.dry else "" print( f"I/O mode: {workers} workers, {producer.num_slots} buffers, sequential writes{dry_str}", file=sys.stderr, ) progress_state = {"written": 0} progress = ProgressDisplay(total_bytes, start_time, progress_state, seed=seed_for_display) if not args.quiet: progress.start() try: producer.start() producer.run(progress_state) finally: producer.stop() progress.stop() elapsed = time.perf_counter() - start_time if not args.quiet: print_summary(producer.written, elapsed, "wrote", seed=seed_for_display) if args.verbose: print(producer.format_stats_report(), file=sys.stderr) def main(): """Main entry point for the CLI with exception handling.""" try: _main() except (KeyboardInterrupt, BrokenPipeError): sys.exit(1) except ValueError as e: print(f"Error: {e}", file=sys.stderr) sys.exit(1)