#!/usr/bin/env python3
"""SystemMonitor Operator Script — Generate Engagement Report.

Pulls data from SQLite databases synced from the implant and generates
a structured engagement report in both Markdown and HTML formats.

Sections:
    1. Executive Summary
    2. Network Topology & Host Inventory
    3. Captured Credentials
    4. DNS Intelligence
    5. Traffic Analysis
    6. Timeline of Events
    7. Recommendations

Usage:
    python3 generate_report.py <data_dir> [--output <path>] [--title <title>]

Examples:
    python3 generate_report.py ./bb-pull-20240115
    python3 generate_report.py ./bb-pull-20240115 --title "ACME Corp Assessment"
"""

import argparse
import json
import os
import sqlite3
import sys
import time
from collections import Counter, defaultdict
from datetime import datetime, timezone
from pathlib import Path


# ---------------------------------------------------------------------------
# Database helpers
# ---------------------------------------------------------------------------

def connect_db(db_path):
    """Connect to a SQLite database if it exists."""
    if not os.path.isfile(db_path):
        return None
    try:
        conn = sqlite3.connect(db_path)
        conn.row_factory = sqlite3.Row
        return conn
    except Exception:
        return None


def safe_query(conn, sql, params=None):
    """Execute a query, returning empty list on error."""
    if conn is None:
        return []
    try:
        cursor = conn.execute(sql, params or ())
        return cursor.fetchall()
    except Exception:
        return []


# ---------------------------------------------------------------------------
# Data collection
# ---------------------------------------------------------------------------

def collect_credentials(data_dir):
    """Collect credentials from credential database and Responder logs."""
    creds = []

    # From credential DB
    conn = connect_db(os.path.join(data_dir, "databases", "credentials.db"))
    if conn:
        rows = safe_query(conn, """
            SELECT timestamp, source_ip, target_ip, target_service,
                   username, domain, credential_type, hashcat_mode
            FROM credentials
            ORDER BY timestamp
        """)
        for r in rows:
            creds.append({
                "timestamp": r["timestamp"],
                "source_ip": r["source_ip"] or "",
                "target_ip": r["target_ip"] or "",
                "service": r["target_service"] or "",
                "username": r["username"] or "",
                "domain": r["domain"] or "",
                "type": r["credential_type"] or "",
                "hashcat_mode": r["hashcat_mode"],
            })
        conn.close()

    return creds


def collect_hosts(data_dir):
    """Collect host inventory from state database."""
    hosts = []
    conn = connect_db(os.path.join(data_dir, "databases", "state.db"))
    if conn:
        rows = safe_query(conn, """
            SELECT key, value FROM kv_store
            WHERE module = 'host_discovery'
        """)
        for r in rows:
            try:
                host_data = json.loads(r["value"])
                if isinstance(host_data, dict):
                    hosts.append(host_data)
            except (json.JSONDecodeError, TypeError):
                pass
        conn.close()

    # Also check topology DB
    conn = connect_db(os.path.join(data_dir, "databases", "topology.db"))
    if conn:
        rows = safe_query(conn, """
            SELECT ip, mac, hostname, os, vendor, first_seen, last_seen
            FROM hosts
            ORDER BY ip
        """)
        for r in rows:
            hosts.append({
                "ip": r["ip"],
                "mac": r["mac"] or "",
                "hostname": r["hostname"] or "",
                "os": r["os"] or "",
                "vendor": r["vendor"] or "",
                "first_seen": r["first_seen"],
                "last_seen": r["last_seen"],
            })
        conn.close()

    return hosts


def collect_dns_stats(data_dir):
    """Collect DNS query statistics."""
    stats = {"total_queries": 0, "top_domains": [], "top_queriers": [], "doh_count": 0}

    conn = connect_db(os.path.join(data_dir, "databases", "dns_queries.db"))
    if not conn:
        return stats

    # Total queries
    rows = safe_query(conn, "SELECT COUNT(*) as cnt FROM dns_queries")
    stats["total_queries"] = rows[0]["cnt"] if rows else 0

    # DoH detections
    rows = safe_query(conn, "SELECT COUNT(*) as cnt FROM dns_queries WHERE is_doh = 1")
    stats["doh_count"] = rows[0]["cnt"] if rows else 0

    # Top queried domains
    rows = safe_query(conn, """
        SELECT domain, COUNT(*) as cnt
        FROM dns_queries WHERE is_doh = 0
        GROUP BY domain ORDER BY cnt DESC LIMIT 20
    """)
    stats["top_domains"] = [(r["domain"], r["cnt"]) for r in rows]

    # Top querier IPs
    rows = safe_query(conn, """
        SELECT source_ip, COUNT(*) as cnt
        FROM dns_queries WHERE is_doh = 0
        GROUP BY source_ip ORDER BY cnt DESC LIMIT 15
    """)
    stats["top_queriers"] = [(r["source_ip"], r["cnt"]) for r in rows]

    conn.close()
    return stats


def collect_module_status(data_dir):
    """Collect module runtime status."""
    modules = {}
    conn = connect_db(os.path.join(data_dir, "databases", "state.db"))
    if conn:
        rows = safe_query(conn, """
            SELECT module, status, started, updated
            FROM module_status
        """)
        for r in rows:
            modules[r["module"]] = {
                "status": r["status"],
                "started": r["started"],
                "updated": r["updated"],
            }
        conn.close()
    return modules


# ---------------------------------------------------------------------------
# Report generation
# ---------------------------------------------------------------------------

def generate_markdown(data_dir, title, creds, hosts, dns_stats, modules):
    """Generate Markdown report."""
    now = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC")
    lines = []

    def add(text=""):
        lines.append(text)

    # Header
    add(f"# {title}")
    add(f"\n**Generated:** {now}")
    add(f"**Data Source:** `{os.path.abspath(data_dir)}`\n")
    add("---\n")

    # Executive Summary
    add("## 1. Executive Summary\n")
    add(f"- **Hosts Discovered:** {len(hosts)}")
    add(f"- **Credentials Captured:** {len(creds)}")

    cred_types = Counter(c["type"] for c in creds)
    for ctype, count in cred_types.most_common():
        add(f"  - {ctype}: {count}")

    unique_users = len(set(c["username"] for c in creds if c["username"]))
    add(f"- **Unique Users:** {unique_users}")
    add(f"- **DNS Queries Logged:** {dns_stats['total_queries']:,}")
    add(f"- **DoH Blind Spots:** {dns_stats['doh_count']}")
    add(f"- **Modules Active:** {sum(1 for m in modules.values() if m['status'] == 'running')}")
    add("")

    # Host Inventory
    add("## 2. Network Topology & Host Inventory\n")
    if hosts:
        add("| IP | MAC | Hostname | OS | Vendor |")
        add("|---|---|---|---|---|")
        for h in hosts[:100]:  # Cap at 100 for readability
            add(f"| {h.get('ip', '')} | {h.get('mac', '')} | {h.get('hostname', '')} "
                f"| {h.get('os', '')} | {h.get('vendor', '')} |")
        if len(hosts) > 100:
            add(f"\n*({len(hosts)} total hosts — showing first 100)*\n")
    else:
        add("*No host data available.*\n")

    # Credentials
    add("\n## 3. Captured Credentials\n")
    if creds:
        add("| Time | User | Domain | Service | Type |")
        add("|---|---|---|---|---|")
        for c in creds[:50]:
            ts = ""
            if c.get("timestamp"):
                try:
                    ts = datetime.fromtimestamp(c["timestamp"], tz=timezone.utc).strftime("%Y-%m-%d %H:%M")
                except Exception:
                    pass
            add(f"| {ts} | {c['username']} | {c['domain']} | {c['service']} | {c['type']} |")
        if len(creds) > 50:
            add(f"\n*({len(creds)} total credentials — showing first 50)*\n")

        add("\n### Credential Summary by Type\n")
        for ctype, count in cred_types.most_common():
            add(f"- **{ctype}**: {count}")

        add("\n### Unique Users by Service\n")
        service_users = defaultdict(set)
        for c in creds:
            if c["username"]:
                service_users[c["service"]].add(c["username"])
        for svc, users in sorted(service_users.items()):
            add(f"- **{svc}**: {', '.join(sorted(users)[:10])}"
                + (f" (+{len(users)-10} more)" if len(users) > 10 else ""))
    else:
        add("*No credentials captured.*\n")

    # DNS Intelligence
    add("\n## 4. DNS Intelligence\n")
    add(f"- **Total Queries:** {dns_stats['total_queries']:,}")
    add(f"- **DoH Detections (blind spots):** {dns_stats['doh_count']}\n")

    if dns_stats["top_domains"]:
        add("### Top Queried Domains\n")
        add("| Domain | Queries |")
        add("|---|---|")
        for domain, count in dns_stats["top_domains"]:
            add(f"| {domain} | {count:,} |")

    if dns_stats["top_queriers"]:
        add("\n### Top DNS Clients\n")
        add("| IP | Queries |")
        add("|---|---|")
        for ip, count in dns_stats["top_queriers"]:
            add(f"| {ip} | {count:,} |")

    # Module Status
    add("\n## 5. Module Status\n")
    if modules:
        add("| Module | Status | Started |")
        add("|---|---|---|")
        for name, info in sorted(modules.items()):
            started = ""
            if info.get("started"):
                try:
                    started = datetime.fromtimestamp(
                        info["started"], tz=timezone.utc
                    ).strftime("%Y-%m-%d %H:%M")
                except Exception:
                    pass
            add(f"| {name} | {info['status']} | {started} |")
    else:
        add("*No module status data available.*\n")

    # Recommendations
    add("\n## 6. Recommendations\n")
    add("1. **Credential Analysis**: Run `crack_hashes.sh` against captured NTLMv2 hashes")
    add("2. **File Extraction**: Run `extract_files.sh` on PCAPs for document recovery")
    add("3. **Print Jobs**: Check for print traffic with `extract_print_jobs.sh`")
    add("4. **Email Analysis**: Run `extract_emails.sh` for SMTP traffic recovery")
    add("5. **Offline Analysis**: Run Zeek against PCAPs for deep protocol analysis")
    add("")

    add("---\n")
    add(f"*Report generated by SystemMonitor operator tooling — {now}*")

    return "\n".join(lines)


def markdown_to_html(markdown_text, title):
    """Convert Markdown report to standalone HTML."""
    # Simple Markdown-to-HTML conversion (no external deps)
    html_lines = []
    in_table = False
    in_list = False

    html_lines.append(f"""<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>{title}</title>
<style>
body{{font-family:-apple-system,BlinkMacSystemFont,'Segoe UI',Roboto,sans-serif;max-width:1000px;margin:0 auto;padding:40px 20px;color:#1d1d1f;line-height:1.6;background:#fafafa}}
h1{{color:#1a1a2e;border-bottom:3px solid #0066cc;padding-bottom:12px}}
h2{{color:#16213e;margin-top:32px;border-bottom:1px solid #ddd;padding-bottom:8px}}
h3{{color:#333;margin-top:20px}}
table{{border-collapse:collapse;width:100%;margin:12px 0;font-size:14px}}
th,td{{border:1px solid #ddd;padding:8px 12px;text-align:left}}
th{{background:#f0f2f5;font-weight:600}}
tr:nth-child(even){{background:#f9f9f9}}
tr:hover{{background:#f0f2f5}}
code{{background:#f0f2f5;padding:2px 6px;border-radius:4px;font-size:13px}}
hr{{border:none;border-top:1px solid #ddd;margin:24px 0}}
ul{{margin:8px 0}}
li{{margin:4px 0}}
strong{{color:#1a1a2e}}
em{{color:#666}}
</style>
</head>
<body>
""")

    for line in markdown_text.split("\n"):
        stripped = line.strip()

        # Headings
        if stripped.startswith("# "):
            if in_table:
                html_lines.append("</table>")
                in_table = False
            level = len(stripped) - len(stripped.lstrip("#"))
            text = stripped.lstrip("# ").strip()
            html_lines.append(f"<h{level}>{_inline_format(text)}</h{level}>")
            continue

        # Horizontal rule
        if stripped == "---":
            if in_table:
                html_lines.append("</table>")
                in_table = False
            html_lines.append("<hr>")
            continue

        # Table
        if "|" in stripped and stripped.startswith("|"):
            cells = [c.strip() for c in stripped.split("|")[1:-1]]
            if all(c.replace("-", "").replace(":", "") == "" for c in cells):
                continue  # Skip separator row
            if not in_table:
                html_lines.append("<table>")
                in_table = True
                tag = "th"
            else:
                tag = "td"
            row = "".join(f"<{tag}>{_inline_format(c)}</{tag}>" for c in cells)
            html_lines.append(f"<tr>{row}</tr>")
            continue

        if in_table and not stripped.startswith("|"):
            html_lines.append("</table>")
            in_table = False

        # List items
        if stripped.startswith("- ") or stripped.startswith("* "):
            if not in_list:
                html_lines.append("<ul>")
                in_list = True
            text = stripped[2:].strip()
            html_lines.append(f"<li>{_inline_format(text)}</li>")
            continue
        elif stripped.startswith(tuple(f"{i}. " for i in range(1, 20))):
            if not in_list:
                html_lines.append("<ol>")
                in_list = True
            text = stripped.split(". ", 1)[1] if ". " in stripped else stripped
            html_lines.append(f"<li>{_inline_format(text)}</li>")
            continue

        if in_list and not stripped.startswith(("-", "*")) and not stripped[:2].rstrip(".").isdigit():
            html_lines.append("</ul>" if in_list else "</ol>")
            in_list = False

        # Empty line
        if not stripped:
            continue

        # Paragraph
        html_lines.append(f"<p>{_inline_format(stripped)}</p>")

    if in_table:
        html_lines.append("</table>")
    if in_list:
        html_lines.append("</ul>")

    html_lines.append("</body></html>")
    return "\n".join(html_lines)


def _inline_format(text):
    """Apply inline Markdown formatting (bold, code, italic)."""
    import re
    text = re.sub(r"\*\*(.+?)\*\*", r"<strong>\1</strong>", text)
    text = re.sub(r"`(.+?)`", r"<code>\1</code>", text)
    text = re.sub(r"\*(.+?)\*", r"<em>\1</em>", text)
    return text


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description="Generate SystemMonitor engagement report from synced data."
    )
    parser.add_argument("data_dir", help="Path to pulled data directory")
    parser.add_argument("--output", "-o", help="Output directory (default: <data_dir>/report)")
    parser.add_argument("--title", "-t", default="SystemMonitor Engagement Report",
                        help="Report title")
    args = parser.parse_args()

    data_dir = args.data_dir
    output_dir = args.output or os.path.join(data_dir, "report")

    if not os.path.isdir(data_dir):
        print(f"[-] Data directory not found: {data_dir}", file=sys.stderr)
        sys.exit(1)

    os.makedirs(output_dir, exist_ok=True)

    print(f"[*] Generating report: {args.title}")
    print(f"    Data: {os.path.abspath(data_dir)}")
    print(f"    Output: {os.path.abspath(output_dir)}")
    print()

    # Collect data
    print("[*] Collecting credentials...")
    creds = collect_credentials(data_dir)
    print(f"    Found {len(creds)} credentials")

    print("[*] Collecting host inventory...")
    hosts = collect_hosts(data_dir)
    print(f"    Found {len(hosts)} hosts")

    print("[*] Collecting DNS statistics...")
    dns_stats = collect_dns_stats(data_dir)
    print(f"    {dns_stats['total_queries']:,} queries logged")

    print("[*] Collecting module status...")
    modules = collect_module_status(data_dir)
    print(f"    {len(modules)} modules tracked")
    print()

    # Generate Markdown
    print("[*] Generating Markdown report...")
    md_report = generate_markdown(data_dir, args.title, creds, hosts, dns_stats, modules)
    md_path = os.path.join(output_dir, "report.md")
    with open(md_path, "w") as f:
        f.write(md_report)
    print(f"    Saved: {md_path}")

    # Generate HTML
    print("[*] Generating HTML report...")
    html_report = markdown_to_html(md_report, args.title)
    html_path = os.path.join(output_dir, "report.html")
    with open(html_path, "w") as f:
        f.write(html_report)
    print(f"    Saved: {html_path}")

    print()
    print(f"[+] Report generation complete")
    print(f"    Markdown: {md_path}")
    print(f"    HTML:     {html_path}")


if __name__ == "__main__":
    main()