aift_cli

  1#!/usr/bin/env python3
  2"""AIFT CLI — Command-line interface for automated forensic triage.
  3
  4Provides a standalone command-line tool that runs the full AIFT forensic
  5analysis pipeline without starting the Flask web server. Evidence files
  6or folders are processed, artifacts are parsed, AI analysis is performed,
  7and both HTML and JSON reports are generated.
  8
  9Usage:
 10    python aift_cli.py -e /path/to/evidence -p "Investigate suspicious activity"
 11    python aift_cli.py -e /evidence/folder -p @prompt.txt -o /output --profile recommended
 12
 13Attributes:
 14    EXIT_SUCCESS: Exit code for successful completion (0).
 15    EXIT_FAILURE: Exit code for fatal errors (1).
 16    EXIT_PARTIAL: Exit code for partial success with warnings (2).
 17"""
 18
 19from __future__ import annotations
 20
 21import argparse
 22import logging
 23import sys
 24from pathlib import Path
 25from typing import Any
 26
 27from runtime_compat import UnsupportedPythonVersionError, assert_supported_python_version
 28
 29EXIT_SUCCESS = 0
 30EXIT_FAILURE = 1
 31EXIT_PARTIAL = 2
 32
 33AIFT_LOGO = """       d8888 8888888 8888888888 88888888888
 34      d88888   888   888            888
 35     d88P888   888   888            888
 36    d88P 888   888   8888888        888
 37   d88P  888   888   888            888
 38  d88P   888   888   888            888
 39 d8888888888   888   888            888
 40d88P     888 8888888 888            888
 41
 42╭─╮╷   ╭─╴╭─╮╭─╮╭─╴╭╮╷╭─╮╷╭─╴   ╶┬╴╭─╮╷╭─╮╭─╴╭─╴
 43├─┤│   ├╴ │ │├┬╯├╴ │╰┤╰─╮││      │ ├┬╯│├─┤│╶╮├╴
 44╵ ╵╵   ╵  ╰─╯╵╰╴╰─╴╵ ╵╰─╯╵╰─╴    ╵ ╵╰╴╵╵ ╵╰─╯╰─╴"""
 45
 46
 47def _build_parser() -> argparse.ArgumentParser:
 48    """Build and return the argument parser for the AIFT CLI.
 49
 50    Returns:
 51        Configured ``argparse.ArgumentParser`` instance.
 52    """
 53    parser = argparse.ArgumentParser(
 54        prog="aift_cli.py",
 55        description="AIFT - AI Forensic Triage (CLI Mode)",
 56        formatter_class=argparse.RawDescriptionHelpFormatter,
 57    )
 58
 59    required = parser.add_argument_group("required arguments")
 60    required.add_argument(
 61        "-e", "--evidence",
 62        required=True,
 63        metavar="EVIDENCE",
 64        help=(
 65            "Path to evidence file or folder. If a folder is given, all "
 66            "supported evidence files within it will be discovered and processed."
 67        ),
 68    )
 69    required.add_argument(
 70        "-p", "--prompt",
 71        required=True,
 72        metavar="PROMPT",
 73        help=(
 74            "Investigation context / prompt for AI analysis. Can be a string "
 75            "or a path to a text file (prefix with @, e.g., @prompt.txt)."
 76        ),
 77    )
 78
 79    optional = parser.add_argument_group("optional arguments")
 80    optional.add_argument(
 81        "-o", "--output",
 82        metavar="OUTPUT",
 83        default=None,
 84        help=(
 85            "Output directory for reports. If omitted, reports are written to "
 86            "the created case folder under cases/<case_id>/reports."
 87        ),
 88    )
 89    optional.add_argument(
 90        "--profile",
 91        metavar="PROFILE",
 92        default="recommended",
 93        help=(
 94            'Artifact profile name or JSON file path. Defaults to "recommended". '
 95            "Use --list-profiles to see available profiles."
 96        ),
 97    )
 98    optional.add_argument(
 99        "-c", "--config",
100        metavar="CONFIG",
101        default=None,
102        help=(
103            "Path to a YAML config file. Defaults to config/config.yaml "
104            "in the AIFT installation directory."
105        ),
106    )
107    optional.add_argument(
108        "--case-name",
109        metavar="NAME",
110        default=None,
111        help="Human-readable case name. Auto-generated if not provided.",
112    )
113    optional.add_argument(
114        "--skip-hashing",
115        action=argparse.BooleanOptionalAction,
116        default=None,
117        help=(
118            "Skip SHA-256/MD5 hash computation on evidence "
119            "(--no-skip-hashing forces hashing). When neither flag is given, "
120            "the config's evidence.compute_hashes setting decides; hashing "
121            "stays enabled unless that key is set to false."
122        ),
123    )
124    optional.add_argument(
125        "--date-start",
126        metavar="DATE",
127        default=None,
128        help="Start date for analysis filtering (YYYY-MM-DD).",
129    )
130    optional.add_argument(
131        "--date-end",
132        metavar="DATE",
133        default=None,
134        help="End date for analysis filtering (YYYY-MM-DD).",
135    )
136    optional.add_argument(
137        "--quiet",
138        action="store_true",
139        default=False,
140        help="Suppress progress output. Only print final result.",
141    )
142    optional.add_argument(
143        "--no-logo",
144        action="store_true",
145        default=False,
146        help="Skip the startup logo and print only the AIFT version line.",
147    )
148    optional.add_argument(
149        "--verbose",
150        action="store_true",
151        default=False,
152        help="Enable debug-level logging.",
153    )
154    optional.add_argument(
155        "--list-profiles",
156        action="store_true",
157        default=False,
158        help="List available artifact profiles and exit.",
159    )
160    optional.add_argument(
161        "--version",
162        action="store_true",
163        default=False,
164        help="Show version and exit.",
165    )
166
167    return parser
168
169
170def _resolve_prompt(raw_prompt: str) -> str:
171    """Resolve the prompt argument, loading from file if prefixed with ``@``.
172
173    Args:
174        raw_prompt: The raw ``--prompt`` value from the CLI.
175
176    Returns:
177        The investigation prompt string.
178
179    Raises:
180        SystemExit: If the prompt file does not exist or cannot be read.
181    """
182    if raw_prompt.startswith("@"):
183        file_path = Path(raw_prompt[1:]).expanduser().resolve()
184        if not file_path.is_file():
185            print(f"ERROR: Prompt file not found: {file_path}", file=sys.stderr)
186            raise SystemExit(EXIT_FAILURE)
187        try:
188            return file_path.read_text(encoding="utf-8").strip()
189        except Exception as exc:
190            print(f"ERROR: Failed to read prompt file: {exc}", file=sys.stderr)
191            raise SystemExit(EXIT_FAILURE) from None
192    return raw_prompt
193
194
195def _format_duration(seconds: float) -> str:
196    """Format a duration in seconds as a human-readable string.
197
198    Args:
199        seconds: Duration in seconds.
200
201    Returns:
202        Formatted string like ``"2m 05s"`` or ``"45s"``.
203    """
204    minutes = int(seconds) // 60
205    secs = int(seconds) % 60
206    if minutes > 0:
207        return f"{minutes}m {secs:02d}s"
208    return f"{secs}s"
209
210
211def _format_startup_line() -> str:
212    """Return the one-line CLI startup label with centralized version metadata."""
213    from app.utils.version import TOOL_VERSION
214
215    return f"AIFT {TOOL_VERSION} - By Flip Forensics"
216
217
218def _prepare_stdout_for_text(text: str) -> None:
219    """Allow Unicode banner text to be printed on legacy Windows code pages."""
220    encoding = getattr(sys.stdout, "encoding", None) or "utf-8"
221    try:
222        text.encode(encoding)
223    except UnicodeEncodeError:
224        reconfigure = getattr(sys.stdout, "reconfigure", None)
225        if callable(reconfigure):
226            reconfigure(encoding="utf-8")
227
228
229def _print_startup_banner(include_logo: bool = True) -> None:
230    """Print the CLI startup banner.
231
232    Args:
233        include_logo: If True, print the AIFT ASCII logo before the version line.
234            If False, print only the version line.
235    """
236    if include_logo:
237        _prepare_stdout_for_text(AIFT_LOGO)
238        print(AIFT_LOGO)
239        print()
240    print(_format_startup_line())
241
242
243def _make_progress_callback(quiet: bool) -> Any:
244    """Create a progress callback function for the automation engine.
245
246    Args:
247        quiet: If True, return None (suppress progress output).
248
249    Returns:
250        A callback ``(phase, message, percentage) -> None`` or None.
251    """
252    if quiet:
253        return None
254
255    phase_labels = {
256        "discovery": "DISCOVERY",
257        "hashing": "HASHING  ",
258        "parsing": "PARSING  ",
259        "analysis": "ANALYSIS ",
260        "reporting": "REPORTING",
261    }
262
263    def _callback(phase: str, message: str, percentage: float) -> None:
264        """Print a formatted progress line to stdout.
265
266        Args:
267            phase: Pipeline phase name.
268            message: Human-readable status message.
269            percentage: Progress within the phase (0.0--100.0).
270        """
271        label = phase_labels.get(phase, phase.upper().ljust(9))
272        print(f"[{label}] {message}")
273
274    return _callback
275
276
277def _resolve_date_range(
278    date_start: str | None,
279    date_end: str | None,
280) -> tuple[str, str] | None:
281    """Validate CLI date-range arguments and return an engine tuple.
282
283    Args:
284        date_start: Raw ``--date-start`` value.
285        date_end: Raw ``--date-end`` value.
286
287    Returns:
288        ``(start_date, end_date)`` tuple, or ``None`` if no range was supplied.
289
290    Raises:
291        SystemExit: If only one side is provided, or the shared validator
292            rejects the format/range.
293    """
294    if not date_start and not date_end:
295        return None
296    if not date_start or not date_end:
297        print(
298            "ERROR: Both --date-start and --date-end must be provided together.",
299            file=sys.stderr,
300        )
301        raise SystemExit(EXIT_FAILURE)
302
303    from app.utils.artifact_profiles import validate_analysis_date_range
304
305    try:
306        validated = validate_analysis_date_range(
307            {"start_date": date_start, "end_date": date_end}
308        )
309    except ValueError as exc:
310        print(f"ERROR: Invalid date range: {exc}", file=sys.stderr)
311        raise SystemExit(EXIT_FAILURE) from None
312
313    if validated is None:
314        return None
315    return (validated["start_date"], validated["end_date"])
316
317
318def _print_summary(result: Any) -> None:
319    """Print the final summary block after automation completes.
320
321    The evidence line reports how many of the discovered evidence images
322    were processed successfully; images that failed to process are listed
323    in the warnings section.
324
325    Args:
326        result: An ``AutomationResult`` instance from the automation engine.
327    """
328    separator = "=" * 60
329    print()
330    print(separator)
331    if result.success:
332        print("AIFT Automation Complete")
333    else:
334        print("AIFT Automation Complete (with errors)")
335    print(separator)
336    successful_images = getattr(result, "successful_images", 0)
337    print(f"  Case ID:      {result.case_id or 'N/A'}")
338    print(
339        f"  Evidence:     {successful_images} of "
340        f"{len(result.evidence_files)} discovered evidence image(s) "
341        "processed successfully"
342    )
343    print(f"  Duration:     {_format_duration(result.duration_seconds)}")
344    print()
345
346    if result.html_report_path or result.json_report_path:
347        print("  Reports:")
348        if result.html_report_path:
349            print(f"    HTML: {result.html_report_path}")
350        if result.json_report_path:
351            print(f"    JSON: {result.json_report_path}")
352        case_local_html = getattr(result, "case_local_html_report_path", None)
353        if case_local_html and case_local_html != result.html_report_path:
354            print(f"    Case-local HTML: {case_local_html}")
355        case_local_json = getattr(result, "case_local_json_report_path", None)
356        if case_local_json and case_local_json != result.json_report_path:
357            print(f"    Case-local JSON: {case_local_json}")
358        print()
359
360    if getattr(result, "analysis_results_path", None):
361        print("  Case Analysis Payload:")
362        print(f"    analysis_results.json: {result.analysis_results_path}")
363        print()
364
365    if result.errors:
366        print(f"  Errors: {len(result.errors)}")
367        for err in result.errors:
368            print(f"    - {err}")
369        print()
370
371    if result.warnings:
372        print(f"  Warnings: {len(result.warnings)}")
373        for warn in result.warnings:
374            print(f"    - {warn}")
375        print()
376
377    notices = getattr(result, "notices", None) or []
378    if notices:
379        print(f"  Notes: {len(notices)}")
380        for note in notices:
381            print(f"    - {note}")
382        print()
383
384    print(separator)
385
386
387def _list_profiles() -> None:
388    """Load and print all available artifact profiles, then exit.
389
390    Profiles are always listed from the repository ``profile`` directory;
391    the ``--config`` option does not influence profile resolution.
392
393    Raises:
394        SystemExit: Always exits with code 0 after printing.
395    """
396    from app.utils.artifact_profiles import compose_profile_summaries, resolve_profiles_root
397
398    profiles_root = resolve_profiles_root()
399    profiles = compose_profile_summaries(profiles_root)
400
401    if not profiles:
402        print("No artifact profiles found.")
403        raise SystemExit(EXIT_SUCCESS)
404
405    print("Available artifact profiles:\n")
406    for profile in profiles:
407        name = profile.get("name", "unknown")
408        builtin = profile.get("builtin", False)
409        artifact_count = profile.get("artifact_count", 0)
410        tag = " (built-in)" if builtin else ""
411        print(f"  {name}{tag} — {artifact_count} artifacts")
412        notice = str(profile.get("notice", "") or "").strip()
413        if notice:
414            print(f"      Note: {notice}")
415
416    print()
417    raise SystemExit(EXIT_SUCCESS)
418
419
420def _show_version() -> None:
421    """Print the AIFT version and exit.
422
423    Raises:
424        SystemExit: Always exits with code 0 after printing.
425    """
426    from app.utils.version import TOOL_VERSION
427
428    print(f"AIFT v{TOOL_VERSION}")
429    raise SystemExit(EXIT_SUCCESS)
430
431
432def _configure_logging(verbose: bool) -> None:
433    """Configure Python logging for the CLI session.
434
435    Args:
436        verbose: If True, set loggers to DEBUG. Otherwise, only errors are
437            printed.
438    """
439    level = logging.DEBUG if verbose else logging.ERROR
440    logging.basicConfig(
441        level=level,
442        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
443        stream=sys.stderr,
444    )
445    root_logger = logging.getLogger()
446    root_logger.setLevel(level)
447    for handler in root_logger.handlers:
448        handler.setLevel(level)
449    logging.getLogger("app").setLevel(level)
450
451
452def main() -> None:
453    """Parse arguments and run the AIFT automation pipeline.
454
455    This is the CLI entry point. It validates the Python version, parses
456    command-line arguments, resolves the investigation prompt, builds an
457    ``AutomationRequest``, calls ``run_automation()``, and prints the
458    summary.
459
460    Raises:
461        SystemExit: With exit code 0 (success), 1 (failure), or
462            2 (partial success with warnings).
463    """
464    assert_supported_python_version()
465
466    # Check for early-exit flags before full argument parsing, since
467    # --version and --list-profiles should work without -e and -p.
468    if "--version" in sys.argv[1:]:
469        _show_version()
470    if "--list-profiles" in sys.argv[1:]:
471        _configure_logging("--verbose" in sys.argv[1:])
472        _list_profiles()
473
474    parser = _build_parser()
475    args = parser.parse_args()
476
477    _configure_logging(args.verbose)
478    if not args.quiet:
479        _print_startup_banner(include_logo=not args.no_logo)
480
481    # Resolve prompt (may be a file reference).
482    prompt = _resolve_prompt(args.prompt)
483    if not prompt:
484        print("ERROR: Investigation prompt must not be empty.", file=sys.stderr)
485        raise SystemExit(EXIT_FAILURE)
486
487    date_range = _resolve_date_range(args.date_start, args.date_end)
488
489    # Resolve only explicit output directories.  When omitted, the automation
490    # engine resolves reports into the created case's reports directory.
491    output_dir = Path(args.output).resolve() if args.output else None
492
493    # Lazy-import the automation engine to avoid loading Flask.
494    from app.automation.engine import AutomationRequest, run_automation
495
496    request = AutomationRequest(
497        evidence_path=args.evidence,
498        prompt=prompt,
499        output_dir=output_dir,
500        profile_name=args.profile,
501        config_path=args.config,
502        case_name=args.case_name,
503        skip_hashing=args.skip_hashing,
504        date_range=date_range,
505    )
506
507    progress_callback = _make_progress_callback(args.quiet)
508
509    try:
510        result = run_automation(request, progress_callback=progress_callback)
511    except KeyboardInterrupt:
512        print("\nAborted by user.", file=sys.stderr)
513        raise SystemExit(EXIT_FAILURE) from None
514    except Exception as exc:
515        logging.getLogger(__name__).debug("Unhandled exception", exc_info=True)
516        print(f"ERROR: {exc}", file=sys.stderr)
517        if args.verbose:
518            import traceback
519            traceback.print_exc(file=sys.stderr)
520        raise SystemExit(EXIT_FAILURE) from None
521
522    _print_summary(result)
523
524    if not result.success:
525        raise SystemExit(EXIT_FAILURE)
526    if result.warnings:
527        raise SystemExit(EXIT_PARTIAL)
528    raise SystemExit(EXIT_SUCCESS)
529
530
531if __name__ == "__main__":
532    try:
533        main()
534    except UnsupportedPythonVersionError as error:
535        print(str(error), file=sys.stderr)
536        raise SystemExit(EXIT_FAILURE) from None