#!/usr/bin/env python3
"""Active Directory Vulnerability Assessment agent - parses PingCastle XML
reports and performs LDAP checks to assess AD security posture against
common vulnerability categories."""

import argparse
import json
import xml.etree.ElementTree as ET
from collections import Counter
from datetime import datetime
from pathlib import Path

try:
    from ldap3 import Server, Connection, ALL, SUBTREE
except ImportError:
    Connection = None


def parse_pingcastle_report(xml_path: str) -> dict:
    """Parse PingCastle XML health check report."""
    tree = ET.parse(xml_path)
    root = tree.getroot()

    scores = {}
    for score_elem in root.iter("GlobalScore"):
        scores["global"] = int(score_elem.text) if score_elem.text else 0
    for cat in ("StaleObjectsScore", "PrivilegiedGroupScore",
                "TrustScore", "AnomalyScore"):
        elem = root.find(f".//{cat}")
        if elem is not None and elem.text:
            scores[cat.replace("Score", "").lower()] = int(elem.text)

    risks = []
    for rule in root.iter("HealthcheckRiskRule"):
        rationale = rule.find("Rationale")
        category = rule.find("Category")
        points = rule.find("Points")
        risks.append({
            "rule": rule.find("RiskId").text if rule.find("RiskId") is not None else "",
            "category": category.text if category is not None else "",
            "points": int(points.text) if points is not None and points.text else 0,
            "rationale": rationale.text if rationale is not None else "",
        })

    return {"scores": scores, "risks": risks}


def check_password_policy_ldap(server_url: str, username: str, password: str) -> list[dict]:
    """Check domain password policy via LDAP."""
    if Connection is None:
        return [{"error": "ldap3 not installed"}]
    srv = Server(server_url, get_info=ALL, use_ssl=True)
    conn = Connection(srv, user=username, password=password, auto_bind=True)
    base_dn = conn.server.info.other.get("defaultNamingContext", [""])[0]

    conn.search(base_dn, "(objectClass=domain)", search_scope=SUBTREE,
                attributes=["minPwdLength", "lockoutThreshold", "pwdHistoryLength",
                             "maxPwdAge", "minPwdAge"])
    findings = []
    if conn.entries:
        entry = conn.entries[0]
        min_len = int(str(entry.minPwdLength)) if hasattr(entry, "minPwdLength") else 0
        lockout = int(str(entry.lockoutThreshold)) if hasattr(entry, "lockoutThreshold") else 0
        history = int(str(entry.pwdHistoryLength)) if hasattr(entry, "pwdHistoryLength") else 0

        if min_len < 12:
            findings.append({"check": "min_password_length", "value": min_len,
                             "severity": "high", "detail": f"Min length {min_len} < 12"})
        if lockout == 0:
            findings.append({"check": "account_lockout", "value": lockout,
                             "severity": "critical", "detail": "No account lockout policy"})
        if history < 12:
            findings.append({"check": "password_history", "value": history,
                             "severity": "medium", "detail": f"History {history} < 12"})
    conn.unbind()
    return findings


def check_krbtgt_age(server_url: str, username: str, password: str) -> list[dict]:
    """Check krbtgt account password age."""
    if Connection is None:
        return []
    srv = Server(server_url, get_info=ALL, use_ssl=True)
    conn = Connection(srv, user=username, password=password, auto_bind=True)
    base_dn = conn.server.info.other.get("defaultNamingContext", [""])[0]

    conn.search(base_dn, "(&(objectClass=user)(sAMAccountName=krbtgt))",
                search_scope=SUBTREE, attributes=["pwdLastSet"])
    findings = []
    if conn.entries:
        pwd_set = conn.entries[0].pwdLastSet.value
        if pwd_set:
            age_days = (datetime.utcnow() - pwd_set.replace(tzinfo=None)).days
            if age_days > 180:
                findings.append({
                    "check": "krbtgt_password_age",
                    "value": age_days,
                    "severity": "critical",
                    "detail": f"krbtgt password is {age_days} days old (reset recommended every 180 days)",
                })
    conn.unbind()
    return findings


def assess_pingcastle_risks(risks: list[dict]) -> list[dict]:
    """Convert PingCastle risk rules into standardized findings."""
    findings = []
    for risk in risks:
        severity = "critical" if risk["points"] >= 50 else "high" if risk["points"] >= 20 else "medium" if risk["points"] >= 5 else "low"
        findings.append({
            "type": f"pingcastle_{risk['rule']}",
            "severity": severity,
            "category": risk["category"],
            "points": risk["points"],
            "detail": risk["rationale"],
        })
    return findings


def generate_report(xml_path: str = None, server_url: str = None,
                    username: str = None, password: str = None) -> dict:
    """Run all assessments and build consolidated report."""
    findings = []
    scores = {}

    if xml_path:
        pc_data = parse_pingcastle_report(xml_path)
        scores = pc_data["scores"]
        findings.extend(assess_pingcastle_risks(pc_data["risks"]))

    if server_url and username and password:
        findings.extend(check_password_policy_ldap(server_url, username, password))
        findings.extend(check_krbtgt_age(server_url, username, password))

    severity_counts = Counter(f.get("severity", "info") for f in findings)
    return {
        "report": "ad_vulnerability_assessment",
        "generated_at": datetime.utcnow().isoformat() + "Z",
        "pingcastle_scores": scores,
        "total_findings": len(findings),
        "severity_summary": dict(severity_counts),
        "findings": findings,
    }


def main():
    parser = argparse.ArgumentParser(description="AD Vulnerability Assessment Agent")
    parser.add_argument("--pingcastle-xml", help="PingCastle XML report file")
    parser.add_argument("--server", help="LDAP server URL for live checks")
    parser.add_argument("--username", help="Domain username")
    parser.add_argument("--password", help="Password")
    parser.add_argument("--output", help="Output JSON file path")
    args = parser.parse_args()

    if not args.pingcastle_xml and not args.server:
        parser.error("At least --pingcastle-xml or --server is required")

    report = generate_report(args.pingcastle_xml, args.server, args.username, args.password)
    output = json.dumps(report, indent=2)
    if args.output:
        Path(args.output).write_text(output, encoding="utf-8")
        print(f"Report written to {args.output}")
    else:
        print(output)


if __name__ == "__main__":
    main()