#!/usr/bin/env python3
"""Agent for SOC 2 Type II audit preparation, evidence collection, and compliance monitoring."""

import os
import sys
import json
import argparse
from datetime import datetime, timedelta, timezone

import requests

try:
    import boto3
except ImportError:
    boto3 = None

try:
    import yaml
except ImportError:
    yaml = None


# AICPA Trust Services Criteria - Common Criteria (CC1-CC9)
TRUST_SERVICES_CRITERIA = {
    "CC1": {
        "name": "Control Environment",
        "description": "Organization demonstrates commitment to integrity, ethical values, oversight, structure, authority, responsibility, and competence.",
        "controls": {
            "CC1.1": "Demonstrates commitment to integrity and ethical values",
            "CC1.2": "Board exercises oversight responsibility",
            "CC1.3": "Management establishes structure, authority, and responsibility",
            "CC1.4": "Demonstrates commitment to competence",
            "CC1.5": "Enforces accountability",
        },
    },
    "CC2": {
        "name": "Communication and Information",
        "description": "Organization communicates internal and external information necessary to support the functioning of internal control.",
        "controls": {
            "CC2.1": "Obtains or generates relevant quality information",
            "CC2.2": "Internally communicates information including objectives and responsibilities",
            "CC2.3": "Communicates with external parties regarding matters affecting functioning of internal control",
        },
    },
    "CC3": {
        "name": "Risk Assessment",
        "description": "Organization identifies risks to the achievement of objectives and analyzes risks to determine how they should be managed.",
        "controls": {
            "CC3.1": "Specifies objectives with clarity to enable identification of risks",
            "CC3.2": "Identifies risks to achievement of objectives and analyzes them",
            "CC3.3": "Considers potential for fraud in assessing risks",
            "CC3.4": "Identifies and assesses changes that could impact internal control",
        },
    },
    "CC4": {
        "name": "Monitoring Activities",
        "description": "Organization selects, develops, and performs evaluations to ascertain whether controls are present and functioning.",
        "controls": {
            "CC4.1": "Selects, develops, and performs ongoing and separate evaluations",
            "CC4.2": "Evaluates and communicates internal control deficiencies in a timely manner",
        },
    },
    "CC5": {
        "name": "Control Activities",
        "description": "Organization deploys control activities through policies and procedures that put directives into actions.",
        "controls": {
            "CC5.1": "Selects and develops control activities that contribute to mitigation of risks",
            "CC5.2": "Selects and develops general controls over technology",
            "CC5.3": "Deploys control activities through policies and procedures",
        },
    },
    "CC6": {
        "name": "Logical and Physical Access Controls",
        "description": "Organization restricts logical and physical access to information assets.",
        "controls": {
            "CC6.1": "Implements logical access security software, infrastructure, and architectures",
            "CC6.2": "Prior to issuing system credentials, registers and authorizes new users",
            "CC6.3": "Removes access to protected information assets when appropriate",
            "CC6.4": "Restricts physical access to facilities and protected information assets",
            "CC6.5": "Implements controls to prevent and detect unauthorized access",
            "CC6.6": "Manages points of interaction with external systems",
            "CC6.7": "Restricts the transmission, movement, and removal of information",
            "CC6.8": "Implements controls to prevent or detect unauthorized or malicious software",
        },
    },
    "CC7": {
        "name": "System Operations",
        "description": "Organization detects and monitors system components and anomalies that represent events.",
        "controls": {
            "CC7.1": "Detects and monitors system configuration changes and new vulnerabilities",
            "CC7.2": "Monitors system components for anomalies indicative of malicious acts",
            "CC7.3": "Evaluates detected events and determines whether they constitute incidents",
            "CC7.4": "Responds to identified security incidents",
            "CC7.5": "Identifies, develops, and implements activities to recover from incidents",
        },
    },
    "CC8": {
        "name": "Change Management",
        "description": "Organization authorizes, designs, develops, configures, documents, tests, approves, and implements changes.",
        "controls": {
            "CC8.1": "Authorizes, designs, develops, configures, documents, tests, approves, and implements changes to infrastructure and software",
        },
    },
    "CC9": {
        "name": "Risk Mitigation",
        "description": "Organization identifies, selects, and develops risk mitigation activities for risks arising from business disruption and use of vendors.",
        "controls": {
            "CC9.1": "Identifies and assesses risk mitigation activities for risks from business disruptions",
            "CC9.2": "Assesses and manages risks associated with vendors and business partners",
        },
    },
}


def perform_gap_assessment(controls_status):
    """Perform a gap assessment against all Trust Services Criteria CC1-CC9."""
    results = {
        "assessment_date": datetime.now(timezone.utc).isoformat(),
        "summary": {
            "total_controls": 0, "implemented": 0, "partial": 0,
            "not_implemented": 0, "not_assessed": 0,
        },
        "gaps": [],
        "criteria_status": {},
    }

    for cc_id, cc_data in TRUST_SERVICES_CRITERIA.items():
        criteria_result = {"name": cc_data["name"], "controls": {}}
        for ctrl_id, ctrl_desc in cc_data["controls"].items():
            results["summary"]["total_controls"] += 1
            status_info = controls_status.get(ctrl_id, {})
            status = status_info.get("status", "not_assessed")

            criteria_result["controls"][ctrl_id] = {
                "description": ctrl_desc,
                "status": status,
                "evidence": status_info.get("evidence", ""),
                "gap": status_info.get("gap", ""),
                "owner": status_info.get("owner", ""),
                "due_date": status_info.get("due_date", ""),
            }

            if status == "implemented":
                results["summary"]["implemented"] += 1
            elif status == "partial":
                results["summary"]["partial"] += 1
                results["gaps"].append({
                    "control": ctrl_id, "criteria": cc_id,
                    "description": ctrl_desc,
                    "gap": status_info.get("gap", ""),
                    "severity": "medium",
                })
            elif status == "not_implemented":
                results["summary"]["not_implemented"] += 1
                results["gaps"].append({
                    "control": ctrl_id, "criteria": cc_id,
                    "description": ctrl_desc,
                    "gap": status_info.get("gap", ""),
                    "severity": "high",
                })
            else:
                results["summary"]["not_assessed"] += 1

        results["criteria_status"][cc_id] = criteria_result

    total = results["summary"]["total_controls"]
    implemented = results["summary"]["implemented"]
    results["summary"]["readiness_score"] = (
        round((implemented / total) * 100, 1) if total > 0 else 0
    )
    return results


def collect_aws_iam_evidence():
    """Collect AWS IAM evidence for CC6 (Access Controls)."""
    if boto3 is None:
        return {"status": "error", "error": "boto3 not installed"}

    iam = boto3.client("iam")
    evidence = {
        "collected_at": datetime.now(timezone.utc).isoformat(),
        "criteria": "CC6",
        "type": "aws_iam",
    }

    try:
        users = iam.list_users()["Users"]
        user_details = []
        users_without_mfa = []

        for user in users:
            mfa_devices = iam.list_mfa_devices(UserName=user["UserName"])
            has_mfa = len(mfa_devices["MFADevices"]) > 0
            user_info = {
                "username": user["UserName"],
                "created": user["CreateDate"].isoformat(),
                "mfa_enabled": has_mfa,
                "arn": user["Arn"],
            }
            if "PasswordLastUsed" in user:
                user_info["password_last_used"] = user["PasswordLastUsed"].isoformat()
            user_details.append(user_info)
            if not has_mfa:
                users_without_mfa.append(user["UserName"])

        evidence["users"] = {
            "total": len(users),
            "with_mfa": len(users) - len(users_without_mfa),
            "without_mfa": len(users_without_mfa),
            "mfa_compliance_rate": round(
                ((len(users) - len(users_without_mfa)) / max(len(users), 1)) * 100, 1
            ),
            "users_without_mfa": users_without_mfa,
            "details": user_details,
        }
    except Exception as e:
        evidence["users"] = {"error": str(e)}

    try:
        policy = iam.get_account_password_policy()["PasswordPolicy"]
        evidence["password_policy"] = {
            "minimum_length": policy.get("MinimumPasswordLength", 0),
            "require_symbols": policy.get("RequireSymbols", False),
            "require_numbers": policy.get("RequireNumbers", False),
            "require_uppercase": policy.get("RequireUppercaseCharacters", False),
            "require_lowercase": policy.get("RequireLowercaseCharacters", False),
            "max_age_days": policy.get("MaxPasswordAge", 0),
            "password_reuse_prevention": policy.get("PasswordReusePrevention", 0),
        }
    except Exception as e:
        evidence["password_policy"] = {"error": str(e)}

    return evidence


def collect_aws_cloudtrail_evidence():
    """Collect AWS CloudTrail evidence for CC7 (System Operations)."""
    if boto3 is None:
        return {"status": "error", "error": "boto3 not installed"}

    ct = boto3.client("cloudtrail")
    evidence = {
        "collected_at": datetime.now(timezone.utc).isoformat(),
        "criteria": "CC7",
        "type": "aws_cloudtrail",
    }

    try:
        trails = ct.describe_trails()["trailList"]
        trail_details = []
        for trail in trails:
            status = ct.get_trail_status(Name=trail["TrailARN"])
            trail_details.append({
                "name": trail["Name"],
                "arn": trail["TrailARN"],
                "is_logging": status["IsLogging"],
                "multi_region": trail.get("IsMultiRegionTrail", False),
                "log_validation": trail.get("LogFileValidationEnabled", False),
                "s3_bucket": trail.get("S3BucketName", ""),
                "kms_key": trail.get("KmsKeyId", "none"),
            })
        evidence["trails"] = trail_details
        evidence["all_logging"] = all(t["is_logging"] for t in trail_details)
        evidence["all_multi_region"] = all(t["multi_region"] for t in trail_details)
    except Exception as e:
        evidence["trails"] = {"error": str(e)}

    return evidence


def collect_aws_s3_public_access_evidence():
    """Collect AWS S3 public access evidence for CC6 (Access Controls)."""
    if boto3 is None:
        return {"status": "error", "error": "boto3 not installed"}

    s3 = boto3.client("s3")
    evidence = {
        "collected_at": datetime.now(timezone.utc).isoformat(),
        "criteria": "CC6",
        "type": "aws_s3_public_access",
    }

    try:
        buckets = s3.list_buckets()["Buckets"]
        bucket_details = []
        public_buckets = []

        for bucket in buckets:
            name = bucket["Name"]
            try:
                pab = s3.get_public_access_block(Bucket=name)
                cfg = pab["PublicAccessBlockConfiguration"]
                is_public_blocked = all([
                    cfg.get("BlockPublicAcls", False),
                    cfg.get("IgnorePublicAcls", False),
                    cfg.get("BlockPublicPolicy", False),
                    cfg.get("RestrictPublicBuckets", False),
                ])
            except Exception:
                is_public_blocked = False

            try:
                encryption = s3.get_bucket_encryption(Bucket=name)
                has_encryption = True
                enc_algorithm = (
                    encryption["ServerSideEncryptionConfiguration"]["Rules"][0]
                    ["ApplyServerSideEncryptionByDefault"]["SSEAlgorithm"]
                )
            except Exception:
                has_encryption = False
                enc_algorithm = "none"

            bucket_details.append({
                "name": name,
                "created": bucket["CreationDate"].isoformat(),
                "public_access_blocked": is_public_blocked,
                "encrypted": has_encryption,
                "encryption_algorithm": enc_algorithm,
            })
            if not is_public_blocked:
                public_buckets.append(name)

        evidence["buckets"] = {
            "total": len(buckets),
            "public_access_blocked": len(buckets) - len(public_buckets),
            "potentially_public": len(public_buckets),
            "public_bucket_names": public_buckets,
            "details": bucket_details,
        }
    except Exception as e:
        evidence["buckets"] = {"error": str(e)}

    return evidence


def collect_github_change_management_evidence(org, repo, token, since=None):
    """Collect GitHub PR evidence for CC8 (Change Management)."""
    if since is None:
        since = (datetime.now(timezone.utc) - timedelta(days=90)).isoformat()

    headers = {
        "Authorization": "token " + token,
        "Accept": "application/vnd.github.v3+json",
    }
    evidence = {
        "collected_at": datetime.now(timezone.utc).isoformat(),
        "criteria": "CC8",
        "type": "github_change_management",
        "repository": org + "/" + repo,
    }

    try:
        page = 1
        all_prs = []
        while page <= 10:
            resp = requests.get(
                "https://api.github.com/repos/" + org + "/" + repo + "/pulls",
                params={
                    "state": "closed", "base": "main",
                    "per_page": 100, "page": page,
                    "sort": "updated", "direction": "desc",
                },
                headers=headers, timeout=30,
            )
            resp.raise_for_status()
            prs = resp.json()
            if not prs:
                break
            for pr in prs:
                if pr.get("merged_at") and pr["merged_at"] >= since:
                    all_prs.append(pr)
            page += 1

        pr_details = []
        prs_without_approval = []
        for pr in all_prs:
            reviews_resp = requests.get(
                pr["url"] + "/reviews", headers=headers, timeout=15,
            )
            reviews = reviews_resp.json() if reviews_resp.ok else []
            approved = any(r["state"] == "APPROVED" for r in reviews)
            reviewers = list({r["user"]["login"] for r in reviews if "user" in r})

            pr_details.append({
                "number": pr["number"],
                "title": pr["title"],
                "author": pr["user"]["login"],
                "merged_at": pr["merged_at"],
                "approved": approved,
                "reviewers": reviewers,
                "additions": pr.get("additions", 0),
                "deletions": pr.get("deletions", 0),
            })
            if not approved:
                prs_without_approval.append(pr["number"])

        total = len(pr_details)
        approved_count = total - len(prs_without_approval)
        evidence["pull_requests"] = {
            "total_merged": total,
            "with_approval": approved_count,
            "without_approval": len(prs_without_approval),
            "approval_rate": round((approved_count / max(total, 1)) * 100, 1),
            "exceptions": prs_without_approval,
            "details": pr_details,
        }
    except Exception as e:
        evidence["pull_requests"] = {"error": str(e)}

    return evidence


def collect_branch_protection_evidence(org, repo, token):
    """Collect GitHub branch protection evidence for CC8."""
    headers = {
        "Authorization": "token " + token,
        "Accept": "application/vnd.github.v3+json",
    }
    evidence = {
        "collected_at": datetime.now(timezone.utc).isoformat(),
        "criteria": "CC8",
        "type": "github_branch_protection",
    }

    try:
        resp = requests.get(
            "https://api.github.com/repos/" + org + "/" + repo + "/branches/main/protection",
            headers=headers, timeout=15,
        )
        if resp.status_code == 200:
            protection = resp.json()
            pr_reviews = protection.get("required_pull_request_reviews", {})
            evidence["main_branch"] = {
                "protected": True,
                "required_reviews": pr_reviews.get("required_approving_review_count", 0),
                "dismiss_stale_reviews": pr_reviews.get("dismiss_stale_reviews", False),
                "require_code_owner_reviews": pr_reviews.get("require_code_owner_reviews", False),
                "status_checks_required": protection.get("required_status_checks", {}).get("strict", False),
                "enforce_admins": protection.get("enforce_admins", {}).get("enabled", False),
            }
        elif resp.status_code == 404:
            evidence["main_branch"] = {
                "protected": False,
                "gap": "No branch protection configured",
            }
        else:
            evidence["main_branch"] = {"error": "HTTP " + str(resp.status_code)}
    except Exception as e:
        evidence["main_branch"] = {"error": str(e)}

    return evidence


def run_continuous_compliance_checks(aws_enabled=True):
    """Run a suite of continuous compliance checks and return findings."""
    findings = {
        "check_time": datetime.now(timezone.utc).isoformat(),
        "checks": [], "passing": 0, "failing": 0, "errors": 0,
    }
    checks = []

    if aws_enabled and boto3 is not None:
        # CC6.1: MFA enforcement
        try:
            iam = boto3.client("iam")
            users = iam.list_users()["Users"]
            users_without_mfa = []
            for user in users:
                mfa_devices = iam.list_mfa_devices(UserName=user["UserName"])
                if not mfa_devices["MFADevices"]:
                    users_without_mfa.append(user["UserName"])
            checks.append({
                "control": "CC6.1",
                "check": "All IAM users have MFA enabled",
                "passing": len(users_without_mfa) == 0,
                "details": {"total_users": len(users), "without_mfa": users_without_mfa},
            })
        except Exception as e:
            checks.append({
                "control": "CC6.1",
                "check": "All IAM users have MFA enabled",
                "passing": None, "error": str(e),
            })

        # CC6.6: No public S3 buckets
        try:
            s3 = boto3.client("s3")
            buckets = s3.list_buckets()["Buckets"]
            public_buckets = []
            for bucket in buckets:
                try:
                    pab = s3.get_public_access_block(Bucket=bucket["Name"])
                    cfg = pab["PublicAccessBlockConfiguration"]
                    if not all([
                        cfg.get("BlockPublicAcls", False),
                        cfg.get("IgnorePublicAcls", False),
                        cfg.get("BlockPublicPolicy", False),
                        cfg.get("RestrictPublicBuckets", False),
                    ]):
                        public_buckets.append(bucket["Name"])
                except Exception:
                    public_buckets.append(bucket["Name"])
            checks.append({
                "control": "CC6.6",
                "check": "No public S3 buckets",
                "passing": len(public_buckets) == 0,
                "details": {"public_buckets": public_buckets},
            })
        except Exception as e:
            checks.append({
                "control": "CC6.6",
                "check": "No public S3 buckets",
                "passing": None, "error": str(e),
            })

        # CC7.1: CloudTrail logging enabled
        try:
            ct = boto3.client("cloudtrail")
            trails = ct.describe_trails()["trailList"]
            inactive_trails = []
            for trail in trails:
                status = ct.get_trail_status(Name=trail["TrailARN"])
                if not status["IsLogging"]:
                    inactive_trails.append(trail["Name"])
            checks.append({
                "control": "CC7.1",
                "check": "CloudTrail logging enabled on all trails",
                "passing": len(inactive_trails) == 0,
                "details": {"inactive_trails": inactive_trails},
            })
        except Exception as e:
            checks.append({
                "control": "CC7.1",
                "check": "CloudTrail logging enabled on all trails",
                "passing": None, "error": str(e),
            })

    findings["checks"] = checks
    for check in checks:
        if check.get("passing") is True:
            findings["passing"] += 1
        elif check.get("passing") is False:
            findings["failing"] += 1
        else:
            findings["errors"] += 1

    total = findings["passing"] + findings["failing"]
    findings["compliance_score"] = round(
        (findings["passing"] / max(total, 1)) * 100, 1
    )
    return findings


def generate_remediation_plan(gap_assessment):
    """Generate a prioritized remediation plan from gap assessment results."""
    plan = {
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "total_gaps": len(gap_assessment.get("gaps", [])),
        "remediation_items": [],
    }

    severity_priority = {"high": 1, "medium": 2, "low": 3}
    sorted_gaps = sorted(
        gap_assessment.get("gaps", []),
        key=lambda g: severity_priority.get(g.get("severity", "low"), 3),
    )

    remediation_map = {
        "CC1": (["Draft information security policy", "Establish security committee charter", "Create security awareness training", "Document org structure"], "2-4 weeks"),
        "CC2": (["Publish security policy to all staff", "Document system boundaries", "Establish external communication procedures"], "2-4 weeks"),
        "CC3": (["Conduct formal risk assessment", "Create risk register with owners", "Document fraud risk analysis", "Assess change impacts"], "4-8 weeks"),
        "CC4": (["Implement continuous control monitoring", "Establish internal audit program", "Create deficiency tracking process"], "3-6 weeks"),
        "CC5": (["Map policies to procedures", "Deploy technology controls", "Create control activities documentation"], "3-6 weeks"),
        "CC6": (["Enforce MFA across all accounts", "Automate access provisioning/deprovisioning", "Establish quarterly access reviews", "Deploy encryption at rest and in transit"], "2-4 weeks"),
        "CC7": (["Enable CloudTrail/audit logging everywhere", "Deploy SIEM with anomaly alerting", "Implement weekly vulnerability scanning", "Document and test IR procedures"], "3-6 weeks"),
        "CC8": (["Enforce branch protection with required reviews", "Implement CI/CD with automated testing", "Create change advisory board process", "Document emergency change procedures"], "1-3 weeks"),
        "CC9": (["Implement vendor risk management program", "Create BCP and DR plans", "Conduct annual DR testing", "Collect vendor SOC 2 reports"], "4-8 weeks"),
    }

    for i, gap in enumerate(sorted_gaps, 1):
        actions, effort = remediation_map.get(
            gap["criteria"],
            (["Review and implement appropriate controls"], "2-4 weeks"),
        )
        plan["remediation_items"].append({
            "priority": i,
            "control": gap["control"],
            "criteria": gap["criteria"],
            "description": gap["description"],
            "severity": gap.get("severity", "medium"),
            "recommended_actions": actions,
            "estimated_effort": effort,
        })

    return plan


def generate_evidence_manifest(audit_start, audit_end):
    """Generate a manifest of required evidence packages organized by criteria."""
    manifest = {
        "audit_period": {"start": audit_start, "end": audit_end},
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "evidence_packages": {},
    }

    evidence_map = {
        "CC1": [
            {"name": "Code of Conduct Acknowledgments", "source": "HR system", "frequency": "annual"},
            {"name": "Board Meeting Minutes", "source": "Board secretary", "frequency": "quarterly"},
            {"name": "Organizational Chart", "source": "HR system", "frequency": "annual"},
            {"name": "Background Check Policy", "source": "HR system", "frequency": "annual"},
            {"name": "Security Training Completion Records", "source": "LMS", "frequency": "annual"},
        ],
        "CC2": [
            {"name": "Information Security Policy", "source": "Policy repo", "frequency": "annual"},
            {"name": "System Description Document", "source": "Engineering", "frequency": "annual"},
        ],
        "CC3": [
            {"name": "Risk Assessment Report", "source": "GRC platform", "frequency": "annual"},
            {"name": "Risk Register", "source": "GRC platform", "frequency": "quarterly"},
        ],
        "CC4": [
            {"name": "Control Monitoring Dashboard", "source": "Compliance platform", "frequency": "monthly"},
            {"name": "Internal Audit Reports", "source": "Internal audit", "frequency": "annual"},
        ],
        "CC5": [
            {"name": "IT General Controls Matrix", "source": "GRC platform", "frequency": "annual"},
        ],
        "CC6": [
            {"name": "IAM User MFA Status", "source": "AWS IAM / Okta", "frequency": "daily"},
            {"name": "Access Provisioning Tickets", "source": "Jira/ServiceNow", "frequency": "continuous"},
            {"name": "Quarterly Access Reviews", "source": "IAM system", "frequency": "quarterly"},
            {"name": "Terminated User Access Removal", "source": "HR + IAM", "frequency": "continuous"},
            {"name": "S3 Public Access Report", "source": "AWS S3", "frequency": "daily"},
            {"name": "Encryption Configuration", "source": "AWS KMS", "frequency": "daily"},
            {"name": "Password Policy Config", "source": "AWS IAM / Okta", "frequency": "monthly"},
        ],
        "CC7": [
            {"name": "CloudTrail Logging Status", "source": "AWS CloudTrail", "frequency": "daily"},
            {"name": "SIEM Alert Summaries", "source": "SIEM", "frequency": "monthly"},
            {"name": "Vulnerability Scan Reports", "source": "Scanner", "frequency": "weekly"},
            {"name": "Incident Response Reports", "source": "IR team", "frequency": "per-incident"},
            {"name": "GuardDuty Findings", "source": "AWS GuardDuty", "frequency": "daily"},
        ],
        "CC8": [
            {"name": "PR Approval Records", "source": "GitHub", "frequency": "continuous"},
            {"name": "Branch Protection Config", "source": "GitHub", "frequency": "monthly"},
            {"name": "CI/CD Pipeline Config", "source": "GitHub Actions", "frequency": "monthly"},
        ],
        "CC9": [
            {"name": "Vendor Risk Assessments", "source": "GRC platform", "frequency": "annual"},
            {"name": "Business Continuity Plan", "source": "BCP team", "frequency": "annual"},
            {"name": "DR Test Results", "source": "Engineering", "frequency": "annual"},
            {"name": "Vendor SOC 2 Reports", "source": "Vendors", "frequency": "annual"},
        ],
    }

    for cc_id, items in evidence_map.items():
        criteria_name = TRUST_SERVICES_CRITERIA.get(cc_id, {}).get("name", cc_id)
        manifest["evidence_packages"][cc_id] = {
            "criteria_name": criteria_name,
            "evidence_items": items,
            "item_count": len(items),
        }

    manifest["total_evidence_items"] = sum(
        pkg["item_count"] for pkg in manifest["evidence_packages"].values()
    )
    return manifest


def generate_readiness_report(gap_assessment=None, compliance_checks=None):
    """Generate a comprehensive audit readiness report."""
    report = {
        "report_type": "SOC 2 Type II Audit Readiness",
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "overall_readiness": "unknown",
        "sections": {},
    }

    if gap_assessment:
        report["sections"]["gap_assessment"] = {
            "readiness_score": gap_assessment["summary"]["readiness_score"],
            "total_controls": gap_assessment["summary"]["total_controls"],
            "implemented": gap_assessment["summary"]["implemented"],
            "gaps_remaining": len(gap_assessment["gaps"]),
            "high_severity_gaps": len(
                [g for g in gap_assessment["gaps"] if g["severity"] == "high"]
            ),
        }

    if compliance_checks:
        report["sections"]["continuous_compliance"] = {
            "compliance_score": compliance_checks["compliance_score"],
            "checks_passing": compliance_checks["passing"],
            "checks_failing": compliance_checks["failing"],
            "checks_errored": compliance_checks["errors"],
        }

    scores = []
    if gap_assessment:
        scores.append(gap_assessment["summary"]["readiness_score"])
    if compliance_checks:
        scores.append(compliance_checks["compliance_score"])

    if scores:
        avg_score = sum(scores) / len(scores)
        if avg_score >= 90:
            report["overall_readiness"] = "ready"
        elif avg_score >= 70:
            report["overall_readiness"] = "conditionally_ready"
        elif avg_score >= 50:
            report["overall_readiness"] = "needs_work"
        else:
            report["overall_readiness"] = "not_ready"
        report["overall_score"] = round(avg_score, 1)

    return report


def main():
    parser = argparse.ArgumentParser(
        description="SOC 2 Type II Audit Preparation Agent"
    )
    parser.add_argument(
        "--action",
        choices=[
            "gap-assessment", "collect-aws-iam", "collect-aws-cloudtrail",
            "collect-aws-s3", "collect-github-changes", "collect-branch-protection",
            "compliance-checks", "remediation-plan", "evidence-manifest",
            "readiness-report", "list-criteria", "full-assessment",
        ],
        default="list-criteria",
    )
    parser.add_argument("--output", default="soc2_report.json")
    parser.add_argument("--github-org")
    parser.add_argument("--github-repo")
    parser.add_argument("--github-token", default=os.getenv("GITHUB_TOKEN"))
    parser.add_argument("--audit-start", default="2025-04-01")
    parser.add_argument("--audit-end", default="2026-03-31")
    parser.add_argument("--controls-file")
    parser.add_argument("--aws", action="store_true")
    args = parser.parse_args()

    report = {"generated_at": datetime.now(timezone.utc).isoformat(), "action": args.action}

    if args.action == "list-criteria":
        print("[*] AICPA Trust Services Criteria - Common Criteria (CC1-CC9):\n")
        for cc_id, cc_data in TRUST_SERVICES_CRITERIA.items():
            print("  " + cc_id + ": " + cc_data["name"])
            for ctrl_id, ctrl_desc in cc_data["controls"].items():
                print("       " + ctrl_id + ": " + ctrl_desc[:70])
            print()
        report["criteria"] = TRUST_SERVICES_CRITERIA

    elif args.action == "gap-assessment":
        controls_status = {}
        if args.controls_file and os.path.exists(args.controls_file):
            with open(args.controls_file) as f:
                controls_status = json.load(f)
        else:
            print("[!] No controls file provided. Run with --controls-file")
        result = perform_gap_assessment(controls_status)
        report["gap_assessment"] = result
        s = result["summary"]
        print("[+] Readiness Score: " + str(s["readiness_score"]) + "%")
        print("[+] Implemented: " + str(s["implemented"]) + "/" + str(s["total_controls"]))
        print("[+] Gaps: " + str(len(result["gaps"])))

    elif args.action == "collect-aws-iam":
        print("[+] Collecting AWS IAM evidence (CC6)...")
        result = collect_aws_iam_evidence()
        report["aws_iam_evidence"] = result

    elif args.action == "collect-aws-cloudtrail":
        print("[+] Collecting AWS CloudTrail evidence (CC7)...")
        result = collect_aws_cloudtrail_evidence()
        report["aws_cloudtrail_evidence"] = result

    elif args.action == "collect-aws-s3":
        print("[+] Collecting AWS S3 evidence (CC6)...")
        result = collect_aws_s3_public_access_evidence()
        report["aws_s3_evidence"] = result

    elif args.action == "collect-github-changes":
        if not all([args.github_org, args.github_repo, args.github_token]):
            print("[-] --github-org, --github-repo, and --github-token required")
            sys.exit(1)
        print("[+] Collecting GitHub change evidence...")
        result = collect_github_change_management_evidence(
            args.github_org, args.github_repo, args.github_token,
            since=args.audit_start,
        )
        report["github_change_evidence"] = result

    elif args.action == "collect-branch-protection":
        if not all([args.github_org, args.github_repo, args.github_token]):
            print("[-] --github-org, --github-repo, and --github-token required")
            sys.exit(1)
        result = collect_branch_protection_evidence(
            args.github_org, args.github_repo, args.github_token,
        )
        report["branch_protection_evidence"] = result

    elif args.action == "compliance-checks":
        print("[+] Running continuous compliance checks...")
        result = run_continuous_compliance_checks(aws_enabled=args.aws)
        report["compliance_checks"] = result
        print("[+] Compliance Score: " + str(result["compliance_score"]) + "%")
        for check in result["checks"]:
            p = check.get("passing")
            status = "PASS" if p is True else "FAIL" if p is False else "ERROR"
            print("    [" + status + "] " + check["control"] + ": " + check["check"])

    elif args.action == "remediation-plan":
        controls_status = {}
        if args.controls_file and os.path.exists(args.controls_file):
            with open(args.controls_file) as f:
                controls_status = json.load(f)
        gap = perform_gap_assessment(controls_status)
        plan = generate_remediation_plan(gap)
        report["remediation_plan"] = plan
        print("[+] Remediation plan: " + str(plan["total_gaps"]) + " items")
        for item in plan["remediation_items"]:
            print("    " + str(item["priority"]) + ". [" + item["severity"].upper() + "] " + item["control"])

    elif args.action == "evidence-manifest":
        manifest = generate_evidence_manifest(args.audit_start, args.audit_end)
        report["evidence_manifest"] = manifest
        print("[+] Evidence manifest: " + str(manifest["total_evidence_items"]) + " items")
        for cc_id, pkg in manifest["evidence_packages"].items():
            print("    " + cc_id + " (" + pkg["criteria_name"] + "): " + str(pkg["item_count"]) + " items")

    elif args.action == "readiness-report":
        controls_status = {}
        if args.controls_file and os.path.exists(args.controls_file):
            with open(args.controls_file) as f:
                controls_status = json.load(f)
        gap = perform_gap_assessment(controls_status)
        compliance = run_continuous_compliance_checks(aws_enabled=args.aws)
        readiness = generate_readiness_report(gap_assessment=gap, compliance_checks=compliance)
        report["readiness_report"] = readiness
        print("[+] Overall Readiness: " + readiness["overall_readiness"].upper())

    elif args.action == "full-assessment":
        print("[+] Running full SOC 2 Type II assessment...\n")
        controls_status = {}
        if args.controls_file and os.path.exists(args.controls_file):
            with open(args.controls_file) as f:
                controls_status = json.load(f)
        gap = perform_gap_assessment(controls_status)
        report["gap_assessment"] = gap
        print("[+] Gap Assessment - Readiness: " + str(gap["summary"]["readiness_score"]) + "%")

        manifest = generate_evidence_manifest(args.audit_start, args.audit_end)
        report["evidence_manifest"] = manifest
        print("[+] Evidence Manifest: " + str(manifest["total_evidence_items"]) + " items")

        compliance = run_continuous_compliance_checks(aws_enabled=args.aws)
        report["compliance_checks"] = compliance
        print("[+] Compliance Checks: " + str(compliance["compliance_score"]) + "%")

        plan = generate_remediation_plan(gap)
        report["remediation_plan"] = plan
        print("[+] Remediation Items: " + str(plan["total_gaps"]))

        readiness = generate_readiness_report(gap_assessment=gap, compliance_checks=compliance)
        report["readiness_report"] = readiness
        print("\n[+] OVERALL READINESS: " + readiness["overall_readiness"].upper())
        if "overall_score" in readiness:
            print("[+] OVERALL SCORE: " + str(readiness["overall_score"]) + "%")

    with open(args.output, "w") as f:
        json.dump(report, f, indent=2, default=str)
    print("\n[+] Report saved to " + args.output)


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
SOC 2 Type II Audit Preparation Automation

Automates control mapping to Trust Services Criteria, evidence tracking,
readiness assessment, and audit preparation for SOC 2 Type II examinations.
"""

import json
import csv
import os
from datetime import datetime, timedelta
from pathlib import Path
from dataclasses import dataclass, field, asdict
from enum import Enum
from typing import Optional


class TSCCategory(Enum):
    SECURITY = "Security (Common Criteria)"
    AVAILABILITY = "Availability"
    PROCESSING_INTEGRITY = "Processing Integrity"
    CONFIDENTIALITY = "Confidentiality"
    PRIVACY = "Privacy"


class ControlFrequency(Enum):
    CONTINUOUS = "Continuous"
    DAILY = "Daily"
    WEEKLY = "Weekly"
    MONTHLY = "Monthly"
    QUARTERLY = "Quarterly"
    SEMI_ANNUAL = "Semi-Annual"
    ANNUAL = "Annual"
    AS_NEEDED = "As Needed"


class EvidenceStatus(Enum):
    COLLECTED = "Collected"
    PENDING = "Pending"
    MISSING = "Missing"
    NOT_DUE = "Not Due"


class ControlTestResult(Enum):
    NO_EXCEPTION = "No Exception"
    EXCEPTION_NOTED = "Exception Noted"
    NOT_TESTED = "Not Tested"


# Trust Services Criteria - Common Criteria (Security)
TSC_CRITERIA = {
    "CC1.1": {"series": "CC1", "title": "Demonstrates commitment to integrity and ethical values", "category": "Security"},
    "CC1.2": {"series": "CC1", "title": "Board exercises oversight responsibility", "category": "Security"},
    "CC1.3": {"series": "CC1", "title": "Management establishes structures, reporting lines, and authorities", "category": "Security"},
    "CC1.4": {"series": "CC1", "title": "Demonstrates commitment to attract, develop, and retain competent individuals", "category": "Security"},
    "CC1.5": {"series": "CC1", "title": "Holds individuals accountable for internal control responsibilities", "category": "Security"},
    "CC2.1": {"series": "CC2", "title": "Obtains or generates relevant, quality information", "category": "Security"},
    "CC2.2": {"series": "CC2", "title": "Internally communicates information supporting internal control", "category": "Security"},
    "CC2.3": {"series": "CC2", "title": "Communicates with external parties", "category": "Security"},
    "CC3.1": {"series": "CC3", "title": "Specifies objectives with sufficient clarity", "category": "Security"},
    "CC3.2": {"series": "CC3", "title": "Identifies risks to the achievement of objectives", "category": "Security"},
    "CC3.3": {"series": "CC3", "title": "Considers potential for fraud", "category": "Security"},
    "CC3.4": {"series": "CC3", "title": "Identifies and assesses changes that could impact internal control", "category": "Security"},
    "CC4.1": {"series": "CC4", "title": "Selects, develops, and performs ongoing and separate evaluations", "category": "Security"},
    "CC4.2": {"series": "CC4", "title": "Evaluates and communicates internal control deficiencies", "category": "Security"},
    "CC5.1": {"series": "CC5", "title": "Selects and develops control activities", "category": "Security"},
    "CC5.2": {"series": "CC5", "title": "Selects and develops general controls over technology", "category": "Security"},
    "CC5.3": {"series": "CC5", "title": "Deploys through policies that establish expectations and procedures", "category": "Security"},
    "CC6.1": {"series": "CC6", "title": "Logical access security software, infrastructure, and architectures", "category": "Security"},
    "CC6.2": {"series": "CC6", "title": "Prior to credential issuance, registration and authorization", "category": "Security"},
    "CC6.3": {"series": "CC6", "title": "Access removal and modification upon changes", "category": "Security"},
    "CC6.4": {"series": "CC6", "title": "Physical access restrictions to facilities", "category": "Security"},
    "CC6.5": {"series": "CC6", "title": "Discontinuation of physical access", "category": "Security"},
    "CC6.6": {"series": "CC6", "title": "Logical access security against threats from external sources", "category": "Security"},
    "CC6.7": {"series": "CC6", "title": "Restricts transmission, movement, and removal of information", "category": "Security"},
    "CC6.8": {"series": "CC6", "title": "Controls against threats from unauthorized or malicious code", "category": "Security"},
    "CC7.1": {"series": "CC7", "title": "Detection and monitoring procedures for anomalies", "category": "Security"},
    "CC7.2": {"series": "CC7", "title": "Monitors system components for anomalies", "category": "Security"},
    "CC7.3": {"series": "CC7", "title": "Evaluates security events to determine incidents", "category": "Security"},
    "CC7.4": {"series": "CC7", "title": "Responds to identified security incidents", "category": "Security"},
    "CC7.5": {"series": "CC7", "title": "Identifies and remediates vulnerabilities", "category": "Security"},
    "CC8.1": {"series": "CC8", "title": "Authorizes, designs, develops, tests, approves, and implements changes", "category": "Security"},
    "CC9.1": {"series": "CC9", "title": "Identifies and assesses risk mitigation activities", "category": "Security"},
    "CC9.2": {"series": "CC9", "title": "Assesses and manages risks associated with vendors and partners", "category": "Security"},
    # Availability criteria
    "A1.1": {"series": "A1", "title": "Maintains, monitors, and evaluates current processing capacity", "category": "Availability"},
    "A1.2": {"series": "A1", "title": "Environmental protections, software, data backup and recovery", "category": "Availability"},
    "A1.3": {"series": "A1", "title": "Tests recovery plan procedures", "category": "Availability"},
    # Confidentiality criteria
    "C1.1": {"series": "C1", "title": "Identifies and maintains confidential information", "category": "Confidentiality"},
    "C1.2": {"series": "C1", "title": "Disposes of confidential information", "category": "Confidentiality"},
    # Processing Integrity criteria
    "PI1.1": {"series": "PI1", "title": "Obtains or generates relevant, quality information", "category": "Processing Integrity"},
    "PI1.2": {"series": "PI1", "title": "System inputs are complete, accurate, and timely", "category": "Processing Integrity"},
    "PI1.3": {"series": "PI1", "title": "Processing is complete, valid, accurate, timely, and authorized", "category": "Processing Integrity"},
    "PI1.4": {"series": "PI1", "title": "System output is complete, valid, accurate, timely, and authorized", "category": "Processing Integrity"},
    "PI1.5": {"series": "PI1", "title": "Data stored is complete, valid, accurate, timely, and authorized", "category": "Processing Integrity"},
}


@dataclass
class Control:
    control_id: str
    description: str
    tsc_criteria: list
    control_type: str  # Preventive, Detective, Corrective
    frequency: str
    owner: str
    evidence_type: str
    automated: bool = False
    test_result: str = ControlTestResult.NOT_TESTED.value
    exceptions: list = field(default_factory=list)


@dataclass
class EvidenceItem:
    evidence_id: str
    control_id: str
    tsc_criterion: str
    description: str
    period_start: str
    period_end: str
    collection_date: str = ""
    status: str = EvidenceStatus.PENDING.value
    file_reference: str = ""
    notes: str = ""


@dataclass
class ReadinessItem:
    category: str
    item: str
    status: bool
    notes: str = ""
    remediation: str = ""
    due_date: str = ""


class SOC2AuditPrep:
    """Manages SOC 2 Type II audit preparation."""

    def __init__(self, output_dir: str = "./soc2_output",
                 audit_start: str = "", audit_end: str = ""):
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.controls: list[Control] = []
        self.evidence: list[EvidenceItem] = []
        self.audit_start = audit_start or (datetime.now() - timedelta(days=365)).strftime("%Y-%m-%d")
        self.audit_end = audit_end or datetime.now().strftime("%Y-%m-%d")
        self.selected_categories: list[str] = ["Security"]

    def select_tsc_categories(self, categories: list[str]) -> dict:
        """Select applicable TSC categories for the audit."""
        print("\n" + "=" * 70)
        print("TSC CATEGORY SELECTION")
        print("=" * 70)

        self.selected_categories = categories
        applicable_criteria = {}

        for crit_id, crit_info in TSC_CRITERIA.items():
            if crit_info["category"] in categories:
                applicable_criteria[crit_id] = crit_info

        print(f"\n  Selected Categories: {', '.join(categories)}")
        print(f"  Applicable Criteria: {len(applicable_criteria)}")
        print(f"  Audit Period: {self.audit_start} to {self.audit_end}")

        for cat in categories:
            count = sum(1 for c in applicable_criteria.values() if c["category"] == cat)
            print(f"    {cat}: {count} criteria")

        return applicable_criteria

    def create_control_matrix(self, controls: list[dict]) -> list[Control]:
        """Create control matrix mapping controls to TSC criteria."""
        print("\n" + "=" * 70)
        print("CONTROL MATRIX")
        print("=" * 70)

        for ctrl_data in controls:
            ctrl = Control(**ctrl_data)
            self.controls.append(ctrl)
            criteria_str = ", ".join(ctrl.tsc_criteria)
            print(f"\n  [{ctrl.control_id}] {ctrl.description[:60]}...")
            print(f"    TSC: {criteria_str}")
            print(f"    Type: {ctrl.control_type} | Frequency: {ctrl.frequency}")
            print(f"    Owner: {ctrl.owner} | Automated: {ctrl.automated}")

        # Coverage analysis
        covered_criteria = set()
        for ctrl in self.controls:
            covered_criteria.update(ctrl.tsc_criteria)

        applicable_criteria = {
            k for k, v in TSC_CRITERIA.items()
            if v["category"] in self.selected_categories
        }
        uncovered = applicable_criteria - covered_criteria

        print(f"\n  Total Controls: {len(self.controls)}")
        print(f"  Criteria Covered: {len(covered_criteria)} / {len(applicable_criteria)}")
        if uncovered:
            print(f"  GAPS - Uncovered Criteria: {', '.join(sorted(uncovered))}")
        else:
            print(f"  All applicable criteria covered")

        # Save control matrix
        matrix_path = self.output_dir / "control_matrix.json"
        with open(matrix_path, "w") as f:
            json.dump([asdict(c) for c in self.controls], f, indent=2)

        # CSV for auditor
        csv_path = self.output_dir / "control_matrix.csv"
        with open(csv_path, "w", newline="") as f:
            writer = csv.writer(f)
            writer.writerow([
                "Control ID", "Description", "TSC Criteria", "Control Type",
                "Frequency", "Owner", "Evidence Type", "Automated",
                "Test Result", "Exceptions"
            ])
            for ctrl in self.controls:
                writer.writerow([
                    ctrl.control_id,
                    ctrl.description,
                    "; ".join(ctrl.tsc_criteria),
                    ctrl.control_type,
                    ctrl.frequency,
                    ctrl.owner,
                    ctrl.evidence_type,
                    "Yes" if ctrl.automated else "No",
                    ctrl.test_result,
                    "; ".join(ctrl.exceptions) if ctrl.exceptions else "",
                ])

        print(f"  Control Matrix saved to: {matrix_path}")
        return self.controls

    def track_evidence(self, evidence_items: list[dict]) -> dict:
        """Track evidence collection status for audit period."""
        print("\n" + "=" * 70)
        print("EVIDENCE COLLECTION TRACKER")
        print("=" * 70)

        for item_data in evidence_items:
            item = EvidenceItem(**item_data)
            self.evidence.append(item)

        # Status summary
        status_counts = {}
        for item in self.evidence:
            status_counts.setdefault(item.status, 0)
            status_counts[item.status] += 1

        total = len(self.evidence)
        collected = status_counts.get(EvidenceStatus.COLLECTED.value, 0)
        pending = status_counts.get(EvidenceStatus.PENDING.value, 0)
        missing = status_counts.get(EvidenceStatus.MISSING.value, 0)

        print(f"\n  Evidence Items: {total}")
        print(f"  Collected: {collected} ({collected/total*100:.1f}%)" if total > 0 else "")
        print(f"  Pending: {pending}")
        print(f"  Missing: {missing}")

        if missing > 0:
            print(f"\n  ALERT: {missing} evidence items are missing!")
            for item in self.evidence:
                if item.status == EvidenceStatus.MISSING.value:
                    print(f"    - [{item.evidence_id}] {item.description} (Control: {item.control_id})")

        # Save evidence tracker
        tracker_path = self.output_dir / "evidence_tracker.json"
        with open(tracker_path, "w") as f:
            json.dump([asdict(e) for e in self.evidence], f, indent=2)

        print(f"\n  Evidence Tracker saved to: {tracker_path}")
        return status_counts

    def assess_readiness(self) -> dict:
        """Perform SOC 2 Type II audit readiness assessment."""
        print("\n" + "=" * 70)
        print("SOC 2 TYPE II READINESS ASSESSMENT")
        print("=" * 70)

        readiness_checks = [
            ReadinessItem("Documentation", "System Description Document completed", False),
            ReadinessItem("Documentation", "Management Assertion Letter drafted", False),
            ReadinessItem("Documentation", "Control matrix documented and reviewed", bool(self.controls)),
            ReadinessItem("Documentation", "Security policies current and approved", False),
            ReadinessItem("Documentation", "Risk assessment completed within audit period", False),

            ReadinessItem("Controls", "All TSC criteria mapped to at least one control", False),
            ReadinessItem("Controls", "Control owners identified and briefed", False),
            ReadinessItem("Controls", "No control design gaps identified", False),
            ReadinessItem("Controls", "Compensating controls documented for exceptions", False),

            ReadinessItem("Evidence", "Evidence collection covers full audit period", False),
            ReadinessItem("Evidence", "Quarterly access reviews completed", False),
            ReadinessItem("Evidence", "Annual penetration test completed", False),
            ReadinessItem("Evidence", "Security awareness training records available", False),
            ReadinessItem("Evidence", "Change management tickets with approvals", False),
            ReadinessItem("Evidence", "Incident response logs available", False),
            ReadinessItem("Evidence", "Vulnerability scan reports for audit period", False),

            ReadinessItem("Operations", "Background checks completed for new hires", False),
            ReadinessItem("Operations", "MFA enabled for all in-scope systems", False),
            ReadinessItem("Operations", "Encryption at rest and in transit verified", False),
            ReadinessItem("Operations", "Backup and recovery testing documented", False),

            ReadinessItem("Vendor", "Subservice organizations identified", False),
            ReadinessItem("Vendor", "Carve-out or inclusive method determined", False),
            ReadinessItem("Vendor", "CUECs documented", False),
            ReadinessItem("Vendor", "Vendor SOC reports reviewed", False),
        ]

        # Auto-check based on data
        if self.controls:
            readiness_checks[5].status = True  # Controls mapped

        if self.evidence:
            collected = sum(1 for e in self.evidence if e.status == EvidenceStatus.COLLECTED.value)
            if collected == len(self.evidence) and len(self.evidence) > 0:
                readiness_checks[9].status = True

        # Display results
        categories = {}
        for check in readiness_checks:
            categories.setdefault(check.category, [])
            categories[check.category].append(check)

        total_checks = len(readiness_checks)
        passed = sum(1 for c in readiness_checks if c.status)
        pct = (passed / total_checks * 100) if total_checks > 0 else 0

        print(f"\n  Readiness Score: {passed}/{total_checks} ({pct:.1f}%)")

        for cat, checks in categories.items():
            print(f"\n  {cat}:")
            for check in checks:
                icon = "[PASS]" if check.status else "[FAIL]"
                print(f"    {icon} {check.item}")

        # Recommendation
        if pct >= 90:
            print(f"\n  RECOMMENDATION: Ready for audit. Schedule with audit firm.")
        elif pct >= 70:
            print(f"\n  RECOMMENDATION: Address remaining items within 2-4 weeks.")
        else:
            print(f"\n  RECOMMENDATION: Significant gaps remain. Delay audit until addressed.")

        # Save readiness report
        report = {
            "date": datetime.now().isoformat(),
            "audit_period": f"{self.audit_start} to {self.audit_end}",
            "readiness_percentage": pct,
            "checks": [asdict(c) for c in readiness_checks],
        }
        report_path = self.output_dir / "readiness_assessment.json"
        with open(report_path, "w") as f:
            json.dump(report, f, indent=2)

        print(f"\n  Readiness Report saved to: {report_path}")
        return report

    def generate_audit_summary(self) -> dict:
        """Generate overall audit preparation summary dashboard."""
        print("\n" + "=" * 70)
        print("SOC 2 TYPE II AUDIT PREPARATION SUMMARY")
        print("=" * 70)

        summary = {
            "generated": datetime.now().isoformat(),
            "audit_period": f"{self.audit_start} to {self.audit_end}",
            "tsc_categories": self.selected_categories,
            "controls": {
                "total": len(self.controls),
                "preventive": sum(1 for c in self.controls if c.control_type == "Preventive"),
                "detective": sum(1 for c in self.controls if c.control_type == "Detective"),
                "corrective": sum(1 for c in self.controls if c.control_type == "Corrective"),
                "automated": sum(1 for c in self.controls if c.automated),
                "manual": sum(1 for c in self.controls if not c.automated),
            },
            "evidence": {
                "total": len(self.evidence),
                "collected": sum(1 for e in self.evidence if e.status == EvidenceStatus.COLLECTED.value),
                "pending": sum(1 for e in self.evidence if e.status == EvidenceStatus.PENDING.value),
                "missing": sum(1 for e in self.evidence if e.status == EvidenceStatus.MISSING.value),
            },
        }

        print(f"\n  Audit Period: {summary['audit_period']}")
        print(f"  TSC Categories: {', '.join(summary['tsc_categories'])}")
        print(f"\n  Controls: {summary['controls']['total']}")
        print(f"    Preventive: {summary['controls']['preventive']}")
        print(f"    Detective: {summary['controls']['detective']}")
        print(f"    Corrective: {summary['controls']['corrective']}")
        print(f"    Automated: {summary['controls']['automated']}")
        print(f"\n  Evidence: {summary['evidence']['total']}")
        print(f"    Collected: {summary['evidence']['collected']}")
        print(f"    Pending: {summary['evidence']['pending']}")
        print(f"    Missing: {summary['evidence']['missing']}")

        summary_path = self.output_dir / "audit_summary.json"
        with open(summary_path, "w") as f:
            json.dump(summary, f, indent=2)

        print(f"\n  Summary saved to: {summary_path}")
        return summary


def main():
    """Run SOC 2 Type II audit preparation."""
    prep = SOC2AuditPrep(
        audit_start="2024-01-01",
        audit_end="2024-12-31"
    )

    # Select TSC categories
    criteria = prep.select_tsc_categories(["Security", "Availability", "Confidentiality"])

    # Create control matrix
    sample_controls = [
        {
            "control_id": "CTRL-001",
            "description": "Multi-factor authentication is required for all access to production systems and sensitive data",
            "tsc_criteria": ["CC6.1", "CC6.6"],
            "control_type": "Preventive",
            "frequency": ControlFrequency.CONTINUOUS.value,
            "owner": "IT Security Manager",
            "evidence_type": "System configuration screenshot",
            "automated": True,
        },
        {
            "control_id": "CTRL-002",
            "description": "Quarterly user access reviews are performed for all in-scope systems with formal approval",
            "tsc_criteria": ["CC6.2", "CC6.3"],
            "control_type": "Detective",
            "frequency": ControlFrequency.QUARTERLY.value,
            "owner": "IT Manager",
            "evidence_type": "Access review completion report with approvals",
            "automated": False,
        },
        {
            "control_id": "CTRL-003",
            "description": "Security events are monitored 24/7 through SIEM with automated alerting for critical events",
            "tsc_criteria": ["CC7.1", "CC7.2", "CC7.3"],
            "control_type": "Detective",
            "frequency": ControlFrequency.CONTINUOUS.value,
            "owner": "SOC Team Lead",
            "evidence_type": "SIEM dashboard and alert configuration",
            "automated": True,
        },
        {
            "control_id": "CTRL-004",
            "description": "All changes to production systems follow change management process with peer review and management approval",
            "tsc_criteria": ["CC8.1"],
            "control_type": "Preventive",
            "frequency": ControlFrequency.AS_NEEDED.value,
            "owner": "Engineering Manager",
            "evidence_type": "Change tickets with approval chain",
            "automated": False,
        },
        {
            "control_id": "CTRL-005",
            "description": "Annual penetration testing is conducted by qualified third party with findings remediated",
            "tsc_criteria": ["CC7.1", "CC7.5"],
            "control_type": "Detective",
            "frequency": ControlFrequency.ANNUAL.value,
            "owner": "CISO",
            "evidence_type": "Penetration test report and remediation tracker",
            "automated": False,
        },
        {
            "control_id": "CTRL-006",
            "description": "Data at rest is encrypted using AES-256 and data in transit uses TLS 1.2 or higher",
            "tsc_criteria": ["CC6.1", "CC6.7", "C1.1"],
            "control_type": "Preventive",
            "frequency": ControlFrequency.CONTINUOUS.value,
            "owner": "Cloud Architect",
            "evidence_type": "Encryption configuration verification",
            "automated": True,
        },
        {
            "control_id": "CTRL-007",
            "description": "Incident response process is documented, tested annually, and incidents are tracked to resolution",
            "tsc_criteria": ["CC7.3", "CC7.4"],
            "control_type": "Corrective",
            "frequency": ControlFrequency.AS_NEEDED.value,
            "owner": "Security Incident Manager",
            "evidence_type": "IR plan, tabletop exercise report, incident tickets",
            "automated": False,
        },
        {
            "control_id": "CTRL-008",
            "description": "Vendor security assessments are performed prior to engagement and annually thereafter",
            "tsc_criteria": ["CC9.2"],
            "control_type": "Preventive",
            "frequency": ControlFrequency.ANNUAL.value,
            "owner": "Vendor Management",
            "evidence_type": "Vendor assessment questionnaires and SOC reports",
            "automated": False,
        },
    ]
    controls = prep.create_control_matrix(sample_controls)

    # Track evidence
    sample_evidence = [
        {
            "evidence_id": "EVD-001",
            "control_id": "CTRL-001",
            "tsc_criterion": "CC6.1",
            "description": "MFA configuration screenshots for production systems",
            "period_start": "2024-01-01",
            "period_end": "2024-12-31",
            "status": EvidenceStatus.COLLECTED.value,
            "collection_date": "2024-12-15",
            "file_reference": "evidence/CC6.1/mfa_config_2024.pdf",
        },
        {
            "evidence_id": "EVD-002",
            "control_id": "CTRL-002",
            "tsc_criterion": "CC6.3",
            "description": "Q1 2024 User Access Review - All Systems",
            "period_start": "2024-01-01",
            "period_end": "2024-03-31",
            "status": EvidenceStatus.COLLECTED.value,
            "collection_date": "2024-04-05",
            "file_reference": "evidence/CC6.3/q1_access_review.pdf",
        },
        {
            "evidence_id": "EVD-003",
            "control_id": "CTRL-005",
            "tsc_criterion": "CC7.1",
            "description": "Annual penetration test report",
            "period_start": "2024-01-01",
            "period_end": "2024-12-31",
            "status": EvidenceStatus.COLLECTED.value,
            "collection_date": "2024-10-30",
            "file_reference": "evidence/CC7.1/pentest_report_2024.pdf",
        },
        {
            "evidence_id": "EVD-004",
            "control_id": "CTRL-002",
            "tsc_criterion": "CC6.3",
            "description": "Q4 2024 User Access Review - All Systems",
            "period_start": "2024-10-01",
            "period_end": "2024-12-31",
            "status": EvidenceStatus.PENDING.value,
        },
    ]
    prep.track_evidence(sample_evidence)

    # Readiness assessment
    prep.assess_readiness()

    # Generate summary
    prep.generate_audit_summary()

    print("\n" + "=" * 70)
    print("SOC 2 TYPE II AUDIT PREPARATION COMPLETE")
    print("=" * 70)


if __name__ == "__main__":
    main()