#!/usr/bin/env python3
"""Cloud Security Posture Management agent using boto3 for AWS Security Hub and Prowler."""

import argparse
import json
import logging
import os
import subprocess
import sys
from datetime import datetime
from typing import Dict, List, Optional

try:
    import boto3
    from botocore.exceptions import ClientError
except ImportError:
    sys.exit("boto3 required: pip install boto3")

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)


def get_securityhub_client(profile: str = "", region: str = "us-east-1"):
    """Create Security Hub client."""
    session = boto3.Session(profile_name=profile) if profile else boto3.Session()
    return session.client("securityhub", region_name=region)


def get_findings_summary(client, max_results: int = 100) -> dict:
    """Get Security Hub findings grouped by severity."""
    try:
        resp = client.get_findings(
            Filters={"WorkflowStatus": [{"Value": "NEW", "Comparison": "EQUALS"}],
                     "RecordState": [{"Value": "ACTIVE", "Comparison": "EQUALS"}]},
            MaxResults=max_results,
        )
        findings = resp.get("Findings", [])
        by_severity = {"CRITICAL": [], "HIGH": [], "MEDIUM": [], "LOW": []}
        for f in findings:
            sev = f.get("Severity", {}).get("Label", "LOW")
            if sev in by_severity:
                by_severity[sev].append({
                    "title": f.get("Title", ""),
                    "resource": f.get("Resources", [{}])[0].get("Id", ""),
                    "standard": f.get("ProductName", ""),
                    "created": f.get("CreatedAt", ""),
                })
        return {sev: items for sev, items in by_severity.items()}
    except ClientError as exc:
        return {"error": str(exc)}


def get_compliance_summary(client) -> List[dict]:
    """Get compliance status across enabled security standards."""
    try:
        resp = client.get_enabled_standards()
        standards = []
        for sub in resp.get("StandardsSubscriptions", []):
            arn = sub.get("StandardsSubscriptionArn", "")
            controls = client.describe_standards_controls(StandardsSubscriptionArn=arn, MaxResults=100)
            total = len(controls.get("Controls", []))
            passed = sum(1 for c in controls.get("Controls", [])
                         if c.get("ComplianceStatus") == "PASSED")
            standards.append({
                "standard": sub.get("StandardsArn", "").split("/")[-1],
                "status": sub.get("StandardsStatus", ""),
                "total_controls": total,
                "passed": passed,
                "failed": total - passed,
                "compliance_pct": round(passed / total * 100, 1) if total else 0,
            })
        return standards
    except ClientError as exc:
        return [{"error": str(exc)}]


def run_prowler_scan(profile: str = "", region: str = "us-east-1") -> dict:
    """Run Prowler security scan via subprocess."""
    cmd = ["prowler", "aws", "--output-formats", "json", "-M", "json-ocsf"]
    if profile:
        cmd.extend(["--profile", profile])
    if region:
        cmd.extend(["--region", region])
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=600)
        if result.stdout:
            lines = result.stdout.strip().split("\n")
            findings = [json.loads(line) for line in lines if line.strip()]
            return {"findings_count": len(findings), "findings": findings[:50]}
    except (FileNotFoundError, subprocess.TimeoutExpired):
        logger.warning("Prowler not available or timed out")
    return {}


def generate_report(client, profile: str, region: str) -> dict:
    """Generate CSPM report combining Security Hub and Prowler."""
    report = {"analysis_date": datetime.utcnow().isoformat(), "region": region}
    report["security_hub_findings"] = get_findings_summary(client)
    report["compliance_standards"] = get_compliance_summary(client)
    counts = {sev: len(items) for sev, items in report["security_hub_findings"].items()
              if isinstance(items, list)}
    report["summary"] = {
        "finding_counts": counts,
        "total_findings": sum(counts.values()),
        "standards_assessed": len(report["compliance_standards"]),
    }
    report["recommendations"] = []
    if counts.get("CRITICAL", 0) > 0:
        report["recommendations"].append(
            f"Remediate {counts['CRITICAL']} CRITICAL findings immediately")
    return report


def main():
    parser = argparse.ArgumentParser(description="Cloud Security Posture Management Agent")
    parser.add_argument("--profile", default="")
    parser.add_argument("--region", default="us-east-1")
    parser.add_argument("--output-dir", default=".")
    parser.add_argument("--output", default="cspm_report.json")
    args = parser.parse_args()

    os.makedirs(args.output_dir, exist_ok=True)
    client = get_securityhub_client(args.profile, args.region)
    report = generate_report(client, args.profile, args.region)
    out_path = os.path.join(args.output_dir, args.output)
    with open(out_path, "w") as f:
        json.dump(report, f, indent=2, default=str)
    logger.info("Report saved to %s", out_path)
    print(json.dumps(report["summary"], indent=2))


if __name__ == "__main__":
    main()