#!/usr/bin/env python3
"""Build a chronological attack timeline from multiple log and artifact sources.

Reads CSV and JSON source files containing timestamped events (sysmon logs,
firewall logs, file timestamps, email artifacts, PCAP timestamp exports) and
merges them into a unified, sorted timeline.  The output can be rendered as
JSON, Markdown, or CSV.
"""

from __future__ import annotations

import argparse
import csv
import json
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


def _parse_csv_source(path: Path) -> list[dict[str, Any]]:
    """Parse a CSV file into a list of event dicts.

    Expects at least a ``timestamp`` column.  All other columns are preserved
    as event metadata.
    """
    events: list[dict[str, Any]] = []
    with path.open(newline="", encoding="utf-8") as fh:
        reader = csv.DictReader(fh)
        for row in reader:
            events.append({
                "timestamp": row.get("timestamp", ""),
                "source_file": str(path),
                "event_type": row.get("event_type", "unknown"),
                "description": row.get("description", ""),
                "details": {k: v for k, v in row.items()
                            if k not in ("timestamp", "event_type", "description")},
            })
    return events


def _parse_json_source(path: Path) -> list[dict[str, Any]]:
    """Parse a JSON file (list of objects) into event dicts."""
    with path.open(encoding="utf-8") as fh:
        data = json.load(fh)
    if isinstance(data, dict):
        data = data.get("events", [data])
    events: list[dict[str, Any]] = []
    for item in data:
        events.append({
            "timestamp": item.get("timestamp", ""),
            "source_file": str(path),
            "event_type": item.get("event_type", "unknown"),
            "description": item.get("description", ""),
            "details": {k: v for k, v in item.items()
                        if k not in ("timestamp", "event_type", "description")},
        })
    return events


def load_sources(source_paths: list[Path]) -> list[dict[str, Any]]:
    """Load and merge events from all provided source files."""
    all_events: list[dict[str, Any]] = []
    for src in source_paths:
        if not src.exists():
            print(f"[warning] source file not found, skipping: {src}", file=sys.stderr)
            continue
        suffix = src.suffix.lower()
        if suffix == ".csv":
            all_events.extend(_parse_csv_source(src))
        elif suffix == ".json":
            all_events.extend(_parse_json_source(src))
        else:
            print(f"[warning] unsupported file type '{suffix}', skipping: {src}",
                  file=sys.stderr)
    return all_events


def build_timeline(events: list[dict[str, Any]], tz_name: str = "UTC") -> dict[str, Any]:
    """Sort events chronologically and return the complete timeline structure."""
    def _sort_key(evt: dict[str, Any]) -> str:
        return evt.get("timestamp", "")

    sorted_events = sorted(events, key=_sort_key)

    first_ts = sorted_events[0]["timestamp"] if sorted_events else None
    last_ts = sorted_events[-1]["timestamp"] if sorted_events else None

    return {
        "timeline": {
            "generated_utc": datetime.now(timezone.utc).isoformat(),
            "timezone": tz_name,
            "total_events": len(sorted_events),
            "first_event": first_ts,
            "last_event": last_ts,
        },
        "events": sorted_events,
    }


def _render_markdown(timeline: dict[str, Any]) -> str:
    """Render the timeline as a Markdown table."""
    lines: list[str] = [
        "# Attack Timeline",
        "",
        f"- **Total events:** {timeline['timeline']['total_events']}",
        f"- **First event:** {timeline['timeline']['first_event']}",
        f"- **Last event:** {timeline['timeline']['last_event']}",
        f"- **Timezone:** {timeline['timeline']['timezone']}",
        "",
        "| Timestamp | Type | Description | Source |",
        "|-----------|------|-------------|--------|",
    ]
    for evt in timeline["events"]:
        lines.append(
            f"| {evt['timestamp']} | {evt['event_type']} "
            f"| {evt['description']} | {evt['source_file']} |"
        )
    return "\n".join(lines) + "\n"


def _render_csv(timeline: dict[str, Any]) -> str:
    """Render the timeline as CSV text."""
    import io
    buf = io.StringIO()
    writer = csv.writer(buf)
    writer.writerow(["timestamp", "event_type", "description", "source_file"])
    for evt in timeline["events"]:
        writer.writerow([evt["timestamp"], evt["event_type"],
                         evt["description"], evt["source_file"]])
    return buf.getvalue()


def format_output(timeline: dict[str, Any], fmt: str) -> str:
    """Serialize the timeline into the requested format."""
    if fmt == "json":
        return json.dumps(timeline, indent=2)
    elif fmt == "markdown":
        return _render_markdown(timeline)
    elif fmt == "csv":
        return _render_csv(timeline)
    elif fmt == "all":
        parts = [
            "=== JSON ===",
            json.dumps(timeline, indent=2),
            "",
            "=== MARKDOWN ===",
            _render_markdown(timeline),
        ]
        return "\n".join(parts)
    else:
        return json.dumps(timeline, indent=2)


def analyze(input_path: Path, output_format: str = "json") -> dict[str, Any]:
    """Analyze a single source file and return a timeline dict."""
    events = load_sources([input_path])
    return build_timeline(events)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Build a chronological attack timeline from multiple sources."
    )
    parser.add_argument("--input", type=Path, help="Path to a single input source file")
    parser.add_argument("--sources", type=Path, nargs="+",
                        help="Paths to one or more source files (CSV or JSON)")
    parser.add_argument("--output", type=Path,
                        help="Path to output file (stdout if omitted)")
    parser.add_argument("--format", default="json",
                        choices=["json", "markdown", "csv", "all"],
                        help="Output format (default: json)")
    parser.add_argument("--tz", default="UTC",
                        help="Timezone label for the timeline (default: UTC)")
    args = parser.parse_args()

    sources: list[Path] = []
    if args.sources:
        sources.extend(args.sources)
    if args.input:
        sources.append(args.input)

    if not sources:
        parser.error("Provide at least one source via --input or --sources")

    events = load_sources(sources)
    timeline = build_timeline(events, tz_name=args.tz)
    rendered = format_output(timeline, args.format)

    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(rendered, encoding="utf-8")
        print(f"[+] Timeline written to {args.output}", file=sys.stderr)
    else:
        print(rendered)


if __name__ == "__main__":
    main()