#!/usr/bin/env python3 """CLI entry point for the TMT threat modeling toolkit. Run as a script or via the installed 'tmt' console command to execute the full threat modeling loop: pattern scanning, optional LLM review, and report generation. Usage: python run_threat_model.py --target ./src --config config.yaml python run_threat_model.py --target ./src --llm --llm-provider openai --llm-model gpt-4 python run_threat_model.py --target ./src --output-dir ./security-reports """ import argparse import logging import sys from tmt.config import ( TMTConfig, load_config, default_config, LLMConfig, ReportConfig, ScannerConfig, ) from tmt.models import Severity, compute_report_statistics from tmt.runner import ThreatModelRunner # ────────────────────────────────────────────────────────────────────────────── # Configure module-level logging # ────────────────────────────────────────────────────────────────────────────── logging.basicConfig( level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", datefmt="%Y-%m-%d %H:%M:%S", ) logger = logging.getLogger("tmt") def _build_argument_parser() -> argparse.ArgumentParser: """Build the CLI argument parser with all supported options. Returns: Configured ArgumentParser instance. """ parser = argparse.ArgumentParser( prog="tmt", description="TMT - Lightweight Threat Modeling Toolkit for Release Cycles", ) parser.add_argument( "--target", "-t", default=".", help="Target directory to scan (default: current directory)", ) parser.add_argument("--config", "-c", default=None, help="Path to YAML config file") parser.add_argument( "--project-name", "-p", default=None, help="Project name for the report" ) parser.add_argument( "--output-dir", "-o", default="reports", help="Report output directory (default: reports)", ) parser.add_argument( "--formats", nargs="+", default=["markdown", "json"], help="Report formats: markdown json", ) parser.add_argument("--llm", action="store_true", help="Enable LLM-powered review") parser.add_argument( "--llm-provider", default="huggingface", choices=["huggingface", "openai", "anthropic"], help="LLM provider (default: huggingface)", ) parser.add_argument( "--llm-model", default="Qwen/Qwen2.5-72B-Instruct", help="LLM model name (default: Qwen/Qwen2.5-72B-Instruct)", ) return parser def _load_or_build_config(args: argparse.Namespace) -> TMTConfig: """Load config from file or build from CLI arguments. Args: args: Parsed CLI argument namespace. Returns: TMTConfig populated from file or CLI arguments. """ if args.config: return load_config(args.config) return default_config() def _apply_cli_overrides(config: TMTConfig, args: argparse.Namespace) -> TMTConfig: """Apply CLI argument overrides to the loaded configuration. Args: config: Base TMTConfig to modify. args: Parsed CLI argument namespace with overrides. Returns: Modified TMTConfig with CLI overrides applied. """ if args.project_name: config.project_name = args.project_name config.report.output_dir = args.output_dir config.report.formats = args.formats config.llm.enabled = args.llm config.llm.provider = args.llm_provider config.llm.model = args.llm_model return config def _print_summary(report) -> None: """Print a concise findings summary to stdout. Args: report: ThreatModelReport with computed statistics. """ report = compute_report_statistics(report) print(f"\n{'='*60}") print(f" TMT Threat Model Report: {report.project_name}") print(f"{'='*60}") print(f" 🔴 Critical: {report.critical_count}") print(f" 🟠 High: {report.high_count}") print(f" 🟡 Medium: {report.medium_count}") print(f" 🔵 Low: {report.low_count}") print(f" ⚪ Info: {report.info_count}") print(f" ─────────────────────────") print(f" Total: {report.total_findings}") print(f"{'='*60}\n") def _determine_exit_code(report) -> int: """Determine process exit code based on finding severity. Args: report: ThreatModelReport with computed statistics. Returns: Exit code: 2 for critical, 1 for high, 0 otherwise. """ report = compute_report_statistics(report) if report.critical_count > 0: return 2 if report.high_count > 0: return 1 return 0 def main(): """Execute the TMT threat modeling CLI workflow. Parses CLI arguments, loads configuration, runs the threat model loop, prints a summary, and exits with an appropriate code. """ parser = _build_argument_parser() args = parser.parse_args() config = _load_or_build_config(args) config = _apply_cli_overrides(config, args) runner = ThreatModelRunner(config) report = runner.run(target_path=args.target) _print_summary(report) sys.exit(_determine_exit_code(report)) # ────────────────────────────────────────────────────────────────────────────── # Script-level entry: invoke main when executed directly # ────────────────────────────────────────────────────────────────────────────── if __name__ == "__main__": main()