#!/usr/bin/env python3
"""Agent for Active Directory attack path analysis using BloodHound data collection."""

import argparse
import json
import os
import subprocess
from datetime import datetime, timezone


def run_sharphound(domain, username=None, password=None, collection="All"):
    """Execute SharpHound data collection."""
    cmd = ["SharpHound.exe", "-c", collection, "-d", domain]
    if username:
        cmd.extend(["--ldapusername", username])
    if password:
        cmd.extend(["--ldappassword", password])
    try:
        result = subprocess.check_output(cmd, text=True, errors="replace", timeout=120)
        return {"status": "success", "output": result[:500]}
    except (subprocess.SubprocessError, FileNotFoundError):
        return {"status": "failed", "note": "SharpHound.exe not found or execution failed"}


def run_bloodhound_python(domain, username, password, dc_ip, collection="all"):
    """Execute bloodhound-python for cross-platform collection."""
    cmd = [
        "bloodhound-python", "-d", domain, "-u", username, "-p", password,
        "-c", collection, "--zip", "-ns", dc_ip,
    ]
    try:
        result = subprocess.check_output(cmd, text=True, errors="replace", timeout=120)
        return {"status": "success", "output": result[:500]}
    except (subprocess.SubprocessError, FileNotFoundError):
        return {"status": "failed", "note": "bloodhound-python not found"}


def analyze_bloodhound_json(data_dir):
    """Parse BloodHound JSON output for high-value findings."""
    findings = {"users": 0, "computers": 0, "groups": 0, "domains": 0, "attack_paths": []}
    for fname in os.listdir(data_dir):
        fpath = os.path.join(data_dir, fname)
        if not fname.endswith(".json"):
            continue
        try:
            with open(fpath, "r") as f:
                data = json.load(f)
            if "users" in fname.lower():
                users = data.get("data", [])
                findings["users"] = len(users)
                for u in users:
                    props = u.get("Properties", {})
                    if props.get("admincount"):
                        findings["attack_paths"].append({
                            "type": "privileged_user",
                            "name": props.get("name", ""),
                            "enabled": props.get("enabled", False),
                        })
            elif "computers" in fname.lower():
                findings["computers"] = len(data.get("data", []))
            elif "groups" in fname.lower():
                findings["groups"] = len(data.get("data", []))
        except (json.JSONDecodeError, KeyError):
            pass
    return findings


def query_neo4j(query, uri="bolt://localhost:7687", user="neo4j", password="bloodhound"):
    """Execute Cypher query against BloodHound Neo4j database."""
    try:
        from neo4j import GraphDatabase
        driver = GraphDatabase.driver(uri, auth=(user, password))
        with driver.session() as session:
            result = session.run(query)
            records = [dict(r) for r in result]
        driver.close()
        return records
    except ImportError:
        return [{"error": "neo4j driver not installed: pip install neo4j"}]
    except Exception as e:
        return [{"error": str(e)}]


ATTACK_PATH_QUERIES = {
    "shortest_to_da": "MATCH p=shortestPath((u:User {owned:true})-[*1..]->(g:Group {name:'DOMAIN ADMINS@DOMAIN.LOCAL'})) RETURN p",
    "kerberoastable": "MATCH (u:User) WHERE u.hasspn=true AND u.enabled=true RETURN u.name, u.serviceprincipalnames",
    "unconstrained_delegation": "MATCH (c:Computer {unconstraineddelegation:true}) RETURN c.name",
    "dcsync_rights": "MATCH p=(u)-[:GetChanges|GetChangesAll]->(d:Domain) RETURN u.name, d.name",
}


def main():
    parser = argparse.ArgumentParser(
        description="AD attack path analysis with BloodHound (authorized testing only)"
    )
    parser.add_argument("--collect", choices=["sharphound", "bloodhound-python"])
    parser.add_argument("--domain", help="AD domain")
    parser.add_argument("--username", help="Domain username")
    parser.add_argument("--password", help="Domain password")
    parser.add_argument("--dc-ip", help="Domain controller IP")
    parser.add_argument("--analyze-dir", help="Directory with BloodHound JSON files")
    parser.add_argument("--cypher-query", help="Custom Cypher query for Neo4j")
    parser.add_argument("--output", "-o", help="Output JSON report")
    args = parser.parse_args()

    print("[*] BloodHound AD Attack Path Agent")
    print("[!] For authorized security testing only")
    report = {"timestamp": datetime.now(timezone.utc).isoformat(), "findings": {}}

    if args.collect == "sharphound":
        result = run_sharphound(args.domain or "")
        report["findings"]["collection"] = result
    elif args.collect == "bloodhound-python":
        result = run_bloodhound_python(
            args.domain or "", args.username or "", args.password or "", args.dc_ip or ""
        )
        report["findings"]["collection"] = result

    if args.analyze_dir:
        analysis = analyze_bloodhound_json(args.analyze_dir)
        report["findings"]["analysis"] = analysis
        print(f"[*] Users: {analysis['users']}, Computers: {analysis['computers']}")
        print(f"[*] Attack paths found: {len(analysis['attack_paths'])}")

    if args.cypher_query:
        results = query_neo4j(args.cypher_query)
        report["findings"]["cypher_results"] = results

    if args.output:
        with open(args.output, "w") as f:
            json.dump(report, f, indent=2)
        print(f"[*] Report saved to {args.output}")
    else:
        print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()