A-Team-Security-Infra-Agent-Engineer/backend/mcp_cis_server.py

#!/usr/bin/env python3
"""
MCP Server: CIS OCI Foundations Benchmark 3.0 Compliance Scanner

Granular per-section data collection — each section only fetches the OCI data
it needs, avoiding the slow full-tenancy scan.
"""

import asyncio
import concurrent.futures
import datetime
import json
import os
import time
import traceback
from configparser import ConfigParser
from functools import partial
from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import Tool, TextContent

# ──────────────────────────────────────────────────────────────────────
server = Server("cis-compliance")

# session_key -> {checker, sections_collected: set, sections_analyzed: set, ...}
_sessions: dict = {}
SESSION_TTL = 7200  # 2 hours — reduce cold starts
TOOL_TIMEOUT = 1800  # 30 min max per tool call
TOOL_RETRIES = 1     # 1 automatic retry on failure


def _session_key(config_id: str, regions: list[str] | None = None) -> str:
    """Build cache key that includes region filter so different scopes are cached separately."""
    if not regions:
        return config_id
    return f"{config_id}:{','.join(sorted(regions))}"

OCI_CONFIGS_DIR = os.environ.get("OCI_CONFIGS_DIR", "/data/oci_configs")

# ──────────────────────────────────────────────────────────────────────
# Per-section data collectors: which private methods to call
# ──────────────────────────────────────────────────────────────────────

# Section -> (home_region_collectors, regional_collectors, analysis_methods)
SECTION_DEPS = {
    "iam": {
        "home": [
            "_CIS_Report__identity_read_users",
            "_CIS_Report__identity_read_tenancy_password_policy",
            "_CIS_Report__identity_read_dynamic_groups",
            "_CIS_Report__identity_read_tenancy_policies",
        ],
        "regional": [],
        "analyze": [
            "_CIS_Report__cis_check_iam_policies",
            "_CIS_Report__cis_check_password_policies",
            "_CIS_Report__cis_check_users",
            "_CIS_Report__cis_check_dynamic_groups",
        ],
    },
    "networking": {
        "home": [],
        "regional": [
            "_CIS_Report__network_read_network_security_lists",
            "_CIS_Report__network_read_network_security_groups_rules",
            "_CIS_Report__network_read_network_vcns",
            "_CIS_Report__network_read_network_subnets",
            "_CIS_Report__adb_read_adbs",
            "_CIS_Report__oic_read_oics",
            "_CIS_Report__oac_read_oacs",
        ],
        "analyze": ["_CIS_Report__cis_check_network_security"],
    },
    "compute": {
        "home": [],
        "regional": ["_CIS_Report__core_instance_read_compute"],
        "analyze": ["_CIS_Report__cis_check_compute_instances"],
    },
    "logging_monitoring": {
        "home": [
            "_CIS_Report__identity_read_tag_defaults",
            "_CIS_Report__identity_read_tenancy_policies",
        ],
        "regional": [
            "_CIS_Report__events_read_event_rules",
            "_CIS_Report__ons_read_subscriptions",
            "_CIS_Report__logging_read_log_groups_and_logs",
            "_CIS_Report__kms_read_keys",
            "_CIS_Report__os_read_buckets",
            "_CIS_Report__network_read_network_subnets",
            "_CIS_Report__network_read_network_capturefilters",
        ],
        "analyze": ["_CIS_Report__cis_check_tagging_and_monitoring"],
    },
    "storage": {
        "home": [],
        "regional": [
            "_CIS_Report__os_read_buckets",
            "_CIS_Report__block_volume_read_block_volumes",
            "_CIS_Report__boot_volume_read_boot_volumes",
            "_CIS_Report__fss_read_fsss",
        ],
        "analyze": ["_CIS_Report__cis_check_storage"],
    },
    "asset_management": {
        "home": [],
        "regional": ["_CIS_Report__search_resources_in_root_compartment"],
        "analyze": ["_CIS_Report__cis_check_assets"],
    },
}

SECTION_NAMES = {
    "iam": "Identity and Access Management",
    "networking": "Networking",
    "compute": "Compute",
    "logging_monitoring": "Logging and Monitoring",
    "storage": "Storage",
    "asset_management": "Asset Management",
}

SECTION_MAP = {"1": "iam", "2": "networking", "3": "compute", "4": "logging_monitoring", "5": "storage", "6": "asset_management"}


# ──────────────────────────────────────────────────────────────────────
# Helpers
# ──────────────────────────────────────────────────────────────────────

def _serialize_oci_obj(obj, max_depth=3):
    if max_depth <= 0: return str(obj)
    if obj is None: return None
    if isinstance(obj, (str, int, float, bool)): return obj
    if isinstance(obj, datetime.datetime): return obj.isoformat()
    if isinstance(obj, (list, tuple)): return [_serialize_oci_obj(i, max_depth-1) for i in obj]
    if isinstance(obj, dict): return {k: _serialize_oci_obj(v, max_depth-1) for k, v in obj.items()}
    if hasattr(obj, "swagger_types"):
        return {a: _serialize_oci_obj(getattr(obj, a, None), max_depth-1) for a in obj.swagger_types}
    return str(obj)


def _get_config_path(config_id: str) -> str:
    return os.path.join(OCI_CONFIGS_DIR, config_id, "config")


def _list_available_configs() -> list[dict]:
    configs = []
    if not os.path.isdir(OCI_CONFIGS_DIR): return configs
    for entry in sorted(os.listdir(OCI_CONFIGS_DIR)):
        config_path = os.path.join(OCI_CONFIGS_DIR, entry, "config")
        if os.path.isfile(config_path):
            info = {"config_id": entry, "config_path": config_path}
            try:
                cp = ConfigParser(); cp.read(config_path)
                sec = cp.sections()[0] if cp.sections() else "DEFAULT"
                info["tenancy_ocid"] = cp.get(sec, "tenancy", fallback="N/A")
                info["region"] = cp.get(sec, "region", fallback="N/A")
                info["user_ocid"] = cp.get(sec, "user", fallback="N/A")
            except Exception: pass
            configs.append(info)
    return configs


def _get_or_create_checker(config_id: str, regions: list[str] | None = None):
    """Return cached session or create a new CIS_Report instance (fast — no data collection)."""
    key = _session_key(config_id, regions)
    now = time.time()
    if key in _sessions:
        s = _sessions[key]
        if now - s["created_at"] < SESSION_TTL:
            return s
        del _sessions[key]

    import cis_reports

    config_path = _get_config_path(config_id)
    if not os.path.isfile(config_path):
        raise FileNotFoundError(f"OCI config not found: {config_path}")

    cp = ConfigParser(); cp.read(config_path)
    profile = cp.sections()[0] if cp.sections() else "DEFAULT"

    config, signer = cis_reports.create_signer(
        file_location=config_path, config_profile=profile,
        is_instance_principals=False, is_delegation_token=False, is_security_token=False,
    )

    regions_str = ",".join(regions) if regions else None

    checker = cis_reports.CIS_Report(
        config=config, signer=signer, proxy=None, output_bucket=None,
        report_directory="/tmp/cis_reports", report_prefix="",
        report_summary_json=False, print_to_screen="FALSE",
        regions_to_run_in=regions_str, raw_data=False, obp=False,
        redact_output=False, debug=False, all_resources=False,
    )

    session = {
        "checker": checker, "created_at": now, "config_id": config_id,
        "regions": regions, "session_key": key,
        "base_collected": False, "sections_collected": set(), "sections_analyzed": set(),
    }
    _sessions[key] = session
    return session


def _ensure_base(session):
    """Collect compartments, groups, domains — shared prereqs for all sections."""
    if session["base_collected"]:
        return
    checker = session["checker"]
    base_methods = [
        "_CIS_Report__identity_read_compartments",
        "_CIS_Report__cloud_guard_read_cloud_guard_configuration",
        "_CIS_Report__identity_read_domains",
        "_CIS_Report__identity_read_groups",
        "_CIS_Report__identity_read_availability_domains",
    ]
    def _run_base(fn_name):
        try:
            getattr(checker, fn_name)()
        except Exception as e:
            print(f"Warning: {fn_name} failed: {e}")
    with concurrent.futures.ThreadPoolExecutor(max_workers=5) as ex:
        list(ex.map(_run_base, base_methods))
    session["base_collected"] = True


def _collect_section(session, section: str):
    """Collect data for a specific section only."""
    if section in session["sections_collected"]:
        return
    _ensure_base(session)
    checker = session["checker"]
    deps = SECTION_DEPS.get(section)
    if not deps:
        return

    if "collection_errors" not in session:
        session["collection_errors"] = {}
    errors = []

    def _run_method(fn_name):
        try:
            getattr(checker, fn_name)()
        except Exception as e:
            errors.append({"collector": fn_name, "error": str(e)[:200]})
            print(f"Warning: {fn_name} failed: {e}")

    # Home region collectors (parallel)
    if deps["home"]:
        with concurrent.futures.ThreadPoolExecutor(max_workers=len(deps["home"])) as ex:
            list(ex.map(_run_method, deps["home"]))

    # Regional collectors (parallel with thread pool)
    if deps["regional"]:
        with concurrent.futures.ThreadPoolExecutor(max_workers=8) as ex:
            list(ex.map(_run_method, deps["regional"]))

    if errors:
        session["collection_errors"][section] = errors

    session["sections_collected"].add(section)


def _analyze_section(session, section: str):
    """Run CIS analysis for a specific section."""
    if section in session["sections_analyzed"]:
        return
    _collect_section(session, section)
    checker = session["checker"]
    deps = SECTION_DEPS.get(section)
    if not deps:
        return

    # Initialize regional findings data if not set
    if not hasattr(checker, '_CIS_Report__cis_regional_findings_data') or not checker._CIS_Report__cis_regional_findings_data:
        checker._CIS_Report__cis_regional_findings_data = {}
        cis_regional_checks = checker._CIS_Report__cis_regional_checks
        for check in cis_regional_checks:
            checker._CIS_Report__cis_regional_findings_data[check] = {}
            for rk in checker._CIS_Report__regions:
                checker._CIS_Report__cis_regional_findings_data[check][rk] = None

    for method_name in deps["analyze"]:
        try:
            getattr(checker, method_name)()
        except Exception as e:
            print(f"Warning: {method_name} failed: {e}")

    session["sections_analyzed"].add(section)


def _check_id_to_section(check_id: str) -> str:
    prefix = check_id.split(".")[0]
    return SECTION_MAP.get(prefix, "unknown")


def _get_benchmark(checker):
    return checker.cis_foundations_benchmark_3_0


def _section_summary(checker, section_key: str) -> dict:
    benchmark = _get_benchmark(checker)
    checks = []
    passed = failed = 0
    for rec_id, check in benchmark.items():
        if _check_id_to_section(rec_id) != section_key:
            continue
        status = check.get("Status")
        if status is True: passed += 1
        elif status is False: failed += 1
        checks.append({
            "check_id": rec_id, "id": check["id"], "title": check["Title"],
            "status": "PASS" if status is True else ("FAIL" if status is False else "N/A"),
            "level": check.get("Level", "N/A"),
            "findings_count": len(check.get("Findings", [])),
            "findings": _serialize_oci_obj(check.get("Findings", [])[:20]),
        })
    total = passed + failed
    score = round(passed / total * 100, 1) if total > 0 else 0
    return {"section": SECTION_NAMES.get(section_key, section_key), "compliance_score": score,
            "passed": passed, "failed": failed, "checks": checks}


def _finding_in_compartment(finding, compartment_id):
    if isinstance(finding, dict):
        fcomp = finding.get("compartment_id")
        if fcomp == compartment_id:
            return True
        if fcomp is None:
            return True  # resource without compartment (IAM tenancy-level) → include
    return False


def _filter_by_compartment(result, compartment_id):
    passed = failed = 0
    for check in result["checks"]:
        original = check["findings"]
        filtered = [f for f in original if _finding_in_compartment(f, compartment_id)]
        check["findings"] = filtered
        check["findings_count"] = len(filtered)
        if check["status"] == "FAIL" and len(filtered) == 0:
            check["status"] = "PASS"
        if check["status"] == "PASS":
            passed += 1
        elif check["status"] == "FAIL":
            failed += 1
    total = passed + failed
    result["passed"] = passed
    result["failed"] = failed
    result["compliance_score"] = round(passed / total * 100, 1) if total > 0 else 0
    result["compartment_filter"] = compartment_id


# ──────────────────────────────────────────────────────────────────────
# Remediation Data
# ──────────────────────────────────────────────────────────────────────

REMEDIATIONS = {
    "1.1": {"description": "Ensure service level admins are created to manage resources of particular service.", "steps": ["Create groups for each service (e.g., VolumeAdmins, ComputeAdmins)", "Write IAM policies granting each group manage access only to their service"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Identity/policieshow/how-policies-work.htm"},
    "1.2": {"description": "Ensure permissions on all resources are given only to the tenancy administrator group.", "steps": ["Review policies with 'manage all-resources in tenancy'", "Remove policies granting full access to non-Administrators groups"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Identity/policieshow/how-policies-work.htm"},
    "1.3": {"description": "Ensure IAM administrators cannot update tenancy Administrators group.", "steps": ["Add condition 'where target.group.name != Administrators' to IAM admin policies"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Identity/policieshow/common-policies.htm"},
    "1.4": {"description": "Ensure IAM password policy requires minimum length of 14 or greater.", "steps": ["Identity > Authentication Settings > Set minimum password length to 14+"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Identity/Tasks/managingpasswordrules.htm"},
    "1.5": {"description": "Ensure IAM password policy expires passwords within 365 days.", "steps": ["Identity > Authentication Settings > Set expiration to 365 days or less"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Identity/Tasks/managingpasswordrules.htm"},
    "1.6": {"description": "Ensure IAM password policy prevents password reuse.", "steps": ["Identity > Authentication Settings > Enable reuse prevention"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Identity/Tasks/managingpasswordrules.htm"},
    "1.7": {"description": "Ensure MFA is enabled for all users with a console password.", "steps": ["Enable MFA for each user with console access", "Enable MFA enforcement in Authentication Settings"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Identity/Tasks/usingmfa.htm"},
    "1.8": {"description": "Ensure user API keys rotate within 90 days.", "steps": ["Identify API keys older than 90 days", "Generate new keys, delete old ones"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Identity/Tasks/managingcredentials.htm"},
    "1.9": {"description": "Ensure user customer secret keys rotate within 90 days.", "steps": ["Identify secret keys older than 90 days", "Rotate and delete old keys"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Identity/Tasks/managingcredentials.htm"},
    "1.10": {"description": "Ensure user auth tokens rotate within 90 days.", "steps": ["Identify auth tokens older than 90 days", "Generate new, delete old"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Identity/Tasks/managingcredentials.htm"},
    "1.11": {"description": "Ensure user IAM Database Passwords rotate within 90 days.", "steps": ["Identify DB passwords older than 90 days", "Reset passwords"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Identity/Tasks/managingcredentials.htm"},
    "1.12": {"description": "Ensure API keys are not created for tenancy administrator users.", "steps": ["Remove API keys from Administrators group members", "Use instance principals for automation"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Identity/Tasks/managingcredentials.htm"},
    "1.13": {"description": "Ensure all IAM users have a valid email address.", "steps": ["Review users and ensure valid email configured"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Identity/Tasks/managingusers.htm"},
    "1.14": {"description": "Ensure Instance Principal authentication is used for instances, ADB and Functions.", "steps": ["Create dynamic groups", "Write policies for dynamic groups", "Update apps to use instance principals"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Identity/Tasks/callingservicesfrominstances.htm"},
    "1.15": {"description": "Ensure storage admins cannot delete resources they manage.", "steps": ["Add 'where request.permission != *_DELETE' conditions"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Identity/policieshow/common-policies.htm"},
    "1.16": {"description": "Ensure credentials unused for 45+ days are disabled.", "steps": ["Identify inactive users", "Disable accounts and credentials"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Identity/Tasks/managingusers.htm"},
    "1.17": {"description": "Ensure only one active API Key per user.", "steps": ["Review users with multiple active API keys", "Remove extras"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Identity/Tasks/managingcredentials.htm"},
    "2.1": {"description": "No security lists allow ingress from 0.0.0.0/0 to port 22.", "steps": ["Review security lists for SSH rules", "Remove or restrict 0.0.0.0/0 to port 22"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Network/Concepts/securitylists.htm"},
    "2.2": {"description": "No security lists allow ingress from 0.0.0.0/0 to port 3389.", "steps": ["Remove or restrict 0.0.0.0/0 to port 3389"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Network/Concepts/securitylists.htm"},
    "2.3": {"description": "No NSGs allow ingress from 0.0.0.0/0 to port 22.", "steps": ["Review and restrict NSG SSH rules"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Network/Concepts/networksecuritygroups.htm"},
    "2.4": {"description": "No NSGs allow ingress from 0.0.0.0/0 to port 3389.", "steps": ["Review and restrict NSG RDP rules"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Network/Concepts/networksecuritygroups.htm"},
    "2.5": {"description": "Default security list restricts all traffic except ICMP.", "steps": ["Remove permissive rules from default security list", "Allow only ICMP within VCN CIDR"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Network/Concepts/securitylists.htm"},
    "2.6": {"description": "OIC access is restricted to allowed sources.", "steps": ["Configure OIC network access restrictions"], "oci_docs": "https://docs.oracle.com/en-us/iaas/integration/doc/restricting-network-access.html"},
    "2.7": {"description": "OAC access is restricted or deployed within a VCN.", "steps": ["Deploy OAC with private access channel"], "oci_docs": "https://docs.oracle.com/en-us/iaas/analytics-cloud/doc/manage-network-access.html"},
    "2.8": {"description": "ADB access is restricted or within a VCN.", "steps": ["Configure ADB with private endpoint"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Database/Concepts/adbsprivateaccess.htm"},
    "3.1": {"description": "Compute Legacy Metadata endpoint is disabled.", "steps": ["Set 'are-legacy-imds-endpoints-disabled' to true"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Compute/Tasks/gettingmetadata.htm"},
    "3.2": {"description": "Secure Boot is enabled on Compute Instance.", "steps": ["Enable Secure Boot in launch options"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Compute/References/shielded-instances.htm"},
    "3.3": {"description": "In-transit Encryption is enabled on Compute Instance.", "steps": ["Enable in-transit encryption in launch options"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Block/Concepts/intransitencryption.htm"},
    "4.1": {"description": "Default tags are used on resources.", "steps": ["Create tag defaults on root compartment"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Tagging/Tasks/managingtagdefaults.htm"},
    "4.2": {"description": "At least one notification topic and subscription exists.", "steps": ["Create ONS topic + subscription"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Notification/Tasks/managingtopicsandsubscriptions.htm"},
    "4.3": {"description": "Notification for Identity Provider changes.", "steps": ["Create Events rule for IdP CRUD"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Events/Task/managingrules.htm"},
    "4.4": {"description": "Notification for IdP group mapping changes.", "steps": ["Create Events rule for IdP mapping CRUD"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Events/Task/managingrules.htm"},
    "4.5": {"description": "Notification for IAM group changes.", "steps": ["Create Events rule for group CRUD"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Events/Task/managingrules.htm"},
    "4.6": {"description": "Notification for IAM policy changes.", "steps": ["Create Events rule for policy CRUD"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Events/Task/managingrules.htm"},
    "4.7": {"description": "Notification for user changes.", "steps": ["Create Events rule for user CRUD"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Events/Task/managingrules.htm"},
    "4.8": {"description": "Notification for VCN changes.", "steps": ["Create Events rule for VCN CRUD"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Events/Task/managingrules.htm"},
    "4.9": {"description": "Notification for route table changes.", "steps": ["Create Events rule for route table CRUD"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Events/Task/managingrules.htm"},
    "4.10": {"description": "Notification for security list changes.", "steps": ["Create Events rule for SL CRUD"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Events/Task/managingrules.htm"},
    "4.11": {"description": "Notification for NSG changes.", "steps": ["Create Events rule for NSG CRUD"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Events/Task/managingrules.htm"},
    "4.12": {"description": "Notification for network gateway changes.", "steps": ["Create Events rules for DRG/IGW/NGW/SGW CRUD"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Events/Task/managingrules.htm"},
    "4.13": {"description": "VCN flow logging is enabled for all subnets.", "steps": ["Enable VCN flow logs on each subnet"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Network/Concepts/vcn-flow-logs.htm"},
    "4.14": {"description": "Cloud Guard is enabled in root compartment.", "steps": ["Enable Cloud Guard in root compartment"], "oci_docs": "https://docs.oracle.com/en-us/iaas/cloud-guard/using/index.htm"},
    "4.15": {"description": "Notification for Cloud Guard problems.", "steps": ["Create Events rule for CG problems"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Events/Task/managingrules.htm"},
    "4.16": {"description": "CMK is rotated at least annually.", "steps": ["Rotate keys older than 365 days"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/KeyManagement/Tasks/managingkeys.htm"},
    "4.17": {"description": "Write level Object Storage logging enabled.", "steps": ["Enable write logs for each bucket"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Logging/Concepts/service_logs.htm"},
    "4.18": {"description": "Notification for Local OCI User Authentication.", "steps": ["Create Events rule for interactive login"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Events/Task/managingrules.htm"},
    "5.1.1": {"description": "No Object Storage buckets are publicly visible.", "steps": ["Set all buckets to NoPublicAccess"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Object/Tasks/managingbuckets.htm"},
    "5.1.2": {"description": "Object Storage encrypted with CMK.", "steps": ["Assign CMK to each bucket"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Object/Tasks/encryption.htm"},
    "5.1.3": {"description": "Versioning enabled for Object Storage.", "steps": ["Enable versioning on each bucket"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Object/Tasks/usingversioning.htm"},
    "5.2.1": {"description": "Block Volumes encrypted with CMK.", "steps": ["Assign CMK to each block volume"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Block/Concepts/blockvolumeencryption.htm"},
    "5.2.2": {"description": "Boot Volumes encrypted with CMK.", "steps": ["Assign CMK to each boot volume"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Block/Concepts/blockvolumeencryption.htm"},
    "5.3.1": {"description": "File Storage encrypted with CMK.", "steps": ["Assign CMK to each file system"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/File/Concepts/filestorageoverview.htm"},
    "6.1": {"description": "At least one compartment exists.", "steps": ["Create compartments for organizing resources"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Identity/Tasks/managingcompartments.htm"},
    "6.2": {"description": "No resources in root compartment.", "steps": ["Move resources to child compartments"], "oci_docs": "https://docs.oracle.com/en-us/iaas/Content/Identity/Tasks/managingcompartments.htm"},
}

# Static checks catalog
CHECKS_CATALOG = {
    "1.1": ("IAM-1", "Identity and Access Management", "Ensure service level admins are created to manage resources of particular service", 1),
    "1.2": ("IAM-2", "Identity and Access Management", "Ensure permissions on all resources are given only to the tenancy administrator group", 1),
    "1.3": ("IAM-3", "Identity and Access Management", "Ensure IAM administrators cannot update tenancy Administrators group", 1),
    "1.4": ("IAM-4", "Identity and Access Management", "Ensure IAM password policy requires minimum length of 14 or greater", 1),
    "1.5": ("IAM-5", "Identity and Access Management", "Ensure IAM password policy expires passwords within 365 days", 1),
    "1.6": ("IAM-6", "Identity and Access Management", "Ensure IAM password policy prevents password reuse", 1),
    "1.7": ("IAM-7", "Identity and Access Management", "Ensure MFA is enabled for all users with a console password", 1),
    "1.8": ("IAM-8", "Identity and Access Management", "Ensure user API keys rotate within 90 days or less", 1),
    "1.9": ("IAM-9", "Identity and Access Management", "Ensure user customer secret keys rotate within 90 days or less", 1),
    "1.10": ("IAM-10", "Identity and Access Management", "Ensure user auth tokens rotate within 90 days or less", 1),
    "1.11": ("IAM-11", "Identity and Access Management", "Ensure user IAM Database Passwords rotate within 90 days", 1),
    "1.12": ("IAM-12", "Identity and Access Management", "Ensure API keys are not created for tenancy administrator users", 1),
    "1.13": ("IAM-13", "Identity and Access Management", "Ensure all OCI IAM user accounts have a valid and current email address", 1),
    "1.14": ("IAM-14", "Identity and Access Management", "Ensure Instance Principal authentication is used", 1),
    "1.15": ("IAM-15", "Identity and Access Management", "Ensure storage service-level admins cannot delete resources they manage", 2),
    "1.16": ("IAM-16", "Identity and Access Management", "Ensure OCI IAM credentials unused for 45 days or more are disabled", 1),
    "1.17": ("IAM-17", "Identity and Access Management", "Ensure there is only one active API Key for any single OCI IAM user", 1),
    "2.1": ("NTW-1", "Networking", "No security lists allow ingress from 0.0.0.0/0 to port 22", 1),
    "2.2": ("NTW-2", "Networking", "No security lists allow ingress from 0.0.0.0/0 to port 3389", 1),
    "2.3": ("NTW-3", "Networking", "No NSGs allow ingress from 0.0.0.0/0 to port 22", 1),
    "2.4": ("NTW-4", "Networking", "No NSGs allow ingress from 0.0.0.0/0 to port 3389", 1),
    "2.5": ("NTW-5", "Networking", "Default security list restricts all traffic except ICMP", 1),
    "2.6": ("NTW-6", "Networking", "OIC access is restricted to allowed sources", 1),
    "2.7": ("NTW-7", "Networking", "OAC access is restricted or deployed within a VCN", 1),
    "2.8": ("NTW-8", "Networking", "ADB access is restricted or deployed within a VCN", 1),
    "3.1": ("COM-1", "Compute", "Legacy Metadata service endpoint is disabled", 2),
    "3.2": ("COM-2", "Compute", "Secure Boot is enabled on Compute Instance", 2),
    "3.3": ("COM-3", "Compute", "In-transit Encryption is enabled on Compute Instance", 1),
    "4.1": ("LAM-1", "Logging and Monitoring", "Default tags are used on resources", 1),
    "4.2": ("LAM-2", "Logging and Monitoring", "At least one notification topic and subscription", 1),
    "4.3": ("LAM-3", "Logging and Monitoring", "Notification for Identity Provider changes", 1),
    "4.4": ("LAM-4", "Logging and Monitoring", "Notification for IdP group mapping changes", 1),
    "4.5": ("LAM-5", "Logging and Monitoring", "Notification for IAM group changes", 1),
    "4.6": ("LAM-6", "Logging and Monitoring", "Notification for IAM policy changes", 1),
    "4.7": ("LAM-7", "Logging and Monitoring", "Notification for user changes", 1),
    "4.8": ("LAM-8", "Logging and Monitoring", "Notification for VCN changes", 1),
    "4.9": ("LAM-9", "Logging and Monitoring", "Notification for route table changes", 1),
    "4.10": ("LAM-10", "Logging and Monitoring", "Notification for security list changes", 1),
    "4.11": ("LAM-11", "Logging and Monitoring", "Notification for NSG changes", 1),
    "4.12": ("LAM-12", "Logging and Monitoring", "Notification for network gateway changes", 1),
    "4.13": ("LAM-13", "Logging and Monitoring", "VCN flow logging enabled for all subnets", 2),
    "4.14": ("LAM-14", "Logging and Monitoring", "Cloud Guard enabled in root compartment", 1),
    "4.15": ("LAM-15", "Logging and Monitoring", "Notification for Cloud Guard problems", 2),
    "4.16": ("LAM-16", "Logging and Monitoring", "CMK rotated at least annually", 1),
    "4.17": ("LAM-17", "Logging and Monitoring", "Write level Object Storage logging enabled", 2),
    "4.18": ("LAM-18", "Logging and Monitoring", "Notification for Local OCI User Authentication", 1),
    "5.1.1": ("STO-1-1", "Storage - Object Storage", "No Object Storage buckets publicly visible", 1),
    "5.1.2": ("STO-1-2", "Storage - Object Storage", "Object Storage encrypted with CMK", 2),
    "5.1.3": ("STO-1-3", "Storage - Object Storage", "Versioning enabled for Object Storage", 2),
    "5.2.1": ("STO-2-1", "Storage - Block Volumes", "Block Volumes encrypted with CMK", 2),
    "5.2.2": ("STO-2-2", "Storage - Block Volumes", "Boot Volumes encrypted with CMK", 2),
    "5.3.1": ("STO-3-1", "Storage - File Storage Service", "File Storage encrypted with CMK", 2),
    "6.1": ("AM-1", "Asset Management", "At least one compartment exists", 1),
    "6.2": ("AM-2", "Asset Management", "No resources in root compartment", 1),
}


# ──────────────────────────────────────────────────────────────────────
# Tool Definitions
# ──────────────────────────────────────────────────────────────────────

TOOLS = [
    Tool(name="cis_list_configs",
         description="Lista as configurações OCI disponíveis. Não requer parâmetros.",
         inputSchema={"type": "object", "properties": {}, "required": []}),
    Tool(name="cis_list_checks",
         description="Lista todos os 48 checks CIS disponíveis com ID, título, seção e nível. Estático, sem conexão OCI.",
         inputSchema={"type": "object", "properties": {
             "section": {"type": "string", "enum": ["iam", "networking", "compute", "logging_monitoring", "storage", "asset_management"], "description": "Filtrar por seção"},
             "level": {"type": "integer", "enum": [1, 2], "description": "Filtrar por nível CIS"},
         }, "required": []}),
    Tool(name="cis_scan_iam",
         description="Coleta dados IAM e executa checks CIS 1.1-1.17: políticas, usuários, MFA, API keys, passwords, dynamic groups. Rápido (~30s).",
         inputSchema={"type": "object", "properties": {
             "config_id": {"type": "string", "description": "ID da config OCI"},
             "regions": {"type": "array", "items": {"type": "string"}, "description": "Regiões específicas (ex: ['us-ashburn-1']). Vazio = todas."},
             "compartment_id": {"type": "string", "description": "OCID de compartment para filtrar findings. Vazio = todos os compartments."},
         }, "required": ["config_id"]}),
    Tool(name="cis_scan_networking",
         description="Coleta dados de rede e executa checks CIS 2.1-2.8: security lists, NSGs, VCNs, subnets, OIC, OAC, ADB.",
         inputSchema={"type": "object", "properties": {
             "config_id": {"type": "string", "description": "ID da config OCI"},
             "regions": {"type": "array", "items": {"type": "string"}, "description": "Regiões específicas (ex: ['us-ashburn-1']). Vazio = todas."},
             "compartment_id": {"type": "string", "description": "OCID de compartment para filtrar findings. Vazio = todos os compartments."},
         }, "required": ["config_id"]}),
    Tool(name="cis_scan_compute",
         description="Coleta instâncias Compute e executa checks CIS 3.1-3.3: metadata v2, secure boot, in-transit encryption. Rápido.",
         inputSchema={"type": "object", "properties": {
             "config_id": {"type": "string", "description": "ID da config OCI"},
             "regions": {"type": "array", "items": {"type": "string"}, "description": "Regiões específicas (ex: ['us-ashburn-1']). Vazio = todas."},
             "compartment_id": {"type": "string", "description": "OCID de compartment para filtrar findings. Vazio = todos os compartments."},
         }, "required": ["config_id"]}),
    Tool(name="cis_scan_logging_monitoring",
         description="Coleta dados de logging/monitoring e executa checks CIS 4.1-4.18: tags, notificações, events, Cloud Guard, KMS, logs.",
         inputSchema={"type": "object", "properties": {
             "config_id": {"type": "string", "description": "ID da config OCI"},
             "regions": {"type": "array", "items": {"type": "string"}, "description": "Regiões específicas (ex: ['us-ashburn-1']). Vazio = todas."},
             "compartment_id": {"type": "string", "description": "OCID de compartment para filtrar findings. Vazio = todos os compartments."},
         }, "required": ["config_id"]}),
    Tool(name="cis_scan_storage",
         description="Coleta dados de storage e executa checks CIS 5.1.1-5.3.1: buckets, block/boot volumes, file storage, CMK, versioning.",
         inputSchema={"type": "object", "properties": {
             "config_id": {"type": "string", "description": "ID da config OCI"},
             "regions": {"type": "array", "items": {"type": "string"}, "description": "Regiões específicas (ex: ['us-ashburn-1']). Vazio = todas."},
             "compartment_id": {"type": "string", "description": "OCID de compartment para filtrar findings. Vazio = todos os compartments."},
         }, "required": ["config_id"]}),
    Tool(name="cis_scan_asset_management",
         description="Coleta dados de assets e executa checks CIS 6.1-6.2: compartments, recursos no root.",
         inputSchema={"type": "object", "properties": {
             "config_id": {"type": "string", "description": "ID da config OCI"},
             "regions": {"type": "array", "items": {"type": "string"}, "description": "Regiões específicas (ex: ['us-ashburn-1']). Vazio = todas."},
             "compartment_id": {"type": "string", "description": "OCID de compartment para filtrar findings. Vazio = todos os compartments."},
         }, "required": ["config_id"]}),
    Tool(name="cis_get_check",
         description="Retorna findings detalhados de um check CIS individual por ID (ex: '1.7', '2.1'). Requer scan da seção correspondente.",
         inputSchema={"type": "object", "properties": {
             "config_id": {"type": "string", "description": "ID da config OCI"},
             "check_id": {"type": "string", "description": "ID do check CIS (ex: '1.7')"},
             "regions": {"type": "array", "items": {"type": "string"}, "description": "Regiões específicas (ex: ['us-ashburn-1']). Vazio = todas."},
         }, "required": ["config_id", "check_id"]}),
    Tool(name="cis_get_remediation",
         description="Orientação de remediação para um check CIS: descrição, passos, link docs Oracle. Estático, sem conexão OCI.",
         inputSchema={"type": "object", "properties": {
             "check_id": {"type": "string", "description": "ID do check CIS (ex: '1.7')"},
         }, "required": ["check_id"]}),
    Tool(name="cis_get_scan_status",
         description="Status da sessão: quais seções já foram coletadas/analisadas, TTL restante, regiões filtradas.",
         inputSchema={"type": "object", "properties": {
             "config_id": {"type": "string", "description": "ID da config OCI"},
             "regions": {"type": "array", "items": {"type": "string"}, "description": "Regiões da sessão a consultar. Vazio = sessão sem filtro."},
         }, "required": ["config_id"]}),
    Tool(name="cis_invalidate_cache",
         description="Limpa cache de sessão. Próximo scan coletará dados novamente.",
         inputSchema={"type": "object", "properties": {
             "config_id": {"type": "string", "description": "ID da config OCI"},
             "regions": {"type": "array", "items": {"type": "string"}, "description": "Regiões da sessão a limpar. Vazio = sessão sem filtro."},
         }, "required": ["config_id"]}),
]


# ──────────────────────────────────────────────────────────────────────
# Handlers
# ──────────────────────────────────────────────────────────────────────

def _handle_cis_list_configs(args):
    configs = _list_available_configs()
    return {"configs": configs, "total": len(configs)}

def _handle_cis_list_checks(args):
    sf = args.get("section"); lf = args.get("level")
    results = []
    for rec_id, (cid, sec, title, lvl) in CHECKS_CATALOG.items():
        if sf and _check_id_to_section(rec_id) != sf: continue
        if lf and lvl != lf: continue
        results.append({"check_id": rec_id, "id": cid, "section": sec, "title": title, "level": lvl})
    return {"checks": results, "total": len(results)}

def _handle_scan_section(args, section: str):
    config_id = args["config_id"]
    regions = args.get("regions") or None
    compartment_id = args.get("compartment_id") or None
    session = _get_or_create_checker(config_id, regions=regions)
    _analyze_section(session, section)
    result = _section_summary(session["checker"], section)
    if regions:
        result["regions_filter"] = regions
    if compartment_id:
        _filter_by_compartment(result, compartment_id)
    col_errors = session.get("collection_errors", {}).get(section, [])
    if col_errors:
        result["partial_errors"] = col_errors
        result["warning"] = f"{len(col_errors)} coletores falharam — resultado pode estar incompleto"
    return result

def _handle_cis_get_check(args):
    config_id = args["config_id"]; check_id = args["check_id"]
    regions = args.get("regions") or None
    section = _check_id_to_section(check_id)
    if section == "unknown":
        return {"error": f"Check '{check_id}' não encontrado. Use cis_list_checks."}
    session = _get_or_create_checker(config_id, regions=regions)
    _analyze_section(session, section)
    benchmark = _get_benchmark(session["checker"])
    if check_id not in benchmark:
        return {"error": f"Check '{check_id}' não encontrado."}
    check = benchmark[check_id]
    report_data = session["checker"].cis_report_data.get(check_id, {})
    return {
        "check_id": check_id, "id": check["id"], "title": check["Title"], "section": check["section"],
        "status": "PASS" if check["Status"] is True else ("FAIL" if check["Status"] is False else "N/A"),
        "level": check.get("Level", "N/A"),
        "total_resources": len(check.get("Total", [])), "findings_count": len(check.get("Findings", [])),
        "findings": _serialize_oci_obj(check.get("Findings", [])[:30]),
        "description": report_data.get("Description", ""),
        "remediation_guidance": report_data.get("Remediation", ""),
    }

def _handle_cis_get_remediation(args):
    cid = args["check_id"]
    if cid not in REMEDIATIONS:
        return {"error": f"Check '{cid}' não encontrado."}
    return {"check_id": cid, **REMEDIATIONS[cid]}

def _handle_cis_get_scan_status(args):
    config_id = args["config_id"]
    regions = args.get("regions") or None
    key = _session_key(config_id, regions)
    if key not in _sessions:
        return {"config_id": config_id, "regions": regions, "status": "not_initialized", "sections_collected": [], "sections_analyzed": []}
    s = _sessions[key]
    elapsed = time.time() - s["created_at"]
    return {
        "config_id": config_id, "regions": s.get("regions"), "status": "active",
        "base_collected": s["base_collected"],
        "sections_collected": list(s["sections_collected"]),
        "sections_analyzed": list(s["sections_analyzed"]),
        "session_age_seconds": round(elapsed),
        "ttl_remaining_seconds": round(max(0, SESSION_TTL - elapsed)),
    }

def _handle_cis_invalidate_cache(args):
    config_id = args["config_id"]
    regions = args.get("regions") or None
    key = _session_key(config_id, regions)
    if key in _sessions:
        del _sessions[key]
        return {"status": "cleared", "config_id": config_id, "regions": regions}
    return {"status": "not_found", "config_id": config_id, "regions": regions}


TOOL_HANDLERS = {
    "cis_list_configs": _handle_cis_list_configs,
    "cis_list_checks": _handle_cis_list_checks,
    "cis_scan_iam": lambda a: _handle_scan_section(a, "iam"),
    "cis_scan_networking": lambda a: _handle_scan_section(a, "networking"),
    "cis_scan_compute": lambda a: _handle_scan_section(a, "compute"),
    "cis_scan_logging_monitoring": lambda a: _handle_scan_section(a, "logging_monitoring"),
    "cis_scan_storage": lambda a: _handle_scan_section(a, "storage"),
    "cis_scan_asset_management": lambda a: _handle_scan_section(a, "asset_management"),
    "cis_get_check": _handle_cis_get_check,
    "cis_get_remediation": _handle_cis_get_remediation,
    "cis_get_scan_status": _handle_cis_get_scan_status,
    "cis_invalidate_cache": _handle_cis_invalidate_cache,
}

# ──────────────────────────────────────────────────────────────────────
# MCP Interface
# ──────────────────────────────────────────────────────────────────────

@server.list_tools()
async def list_tools() -> list[Tool]:
    return TOOLS

@server.call_tool()
async def call_tool(name: str, arguments: dict) -> list[TextContent]:
    handler = TOOL_HANDLERS.get(name)
    if not handler:
        return [TextContent(type="text", text=json.dumps({"error": f"Tool '{name}' not found"}))]
    last_error = None
    loop = asyncio.get_event_loop()
    for attempt in range(1 + TOOL_RETRIES):
        try:
            result = await asyncio.wait_for(
                loop.run_in_executor(None, partial(handler, arguments)),
                timeout=TOOL_TIMEOUT
            )
            return [TextContent(type="text", text=json.dumps(result, default=str, ensure_ascii=False))]
        except asyncio.TimeoutError:
            last_error = f"Tool '{name}' timed out after {TOOL_TIMEOUT}s"
        except Exception as e:
            last_error = f"{e}"
        if attempt < TOOL_RETRIES:
            print(f"[Retry] Tool '{name}' failed ({last_error}), retrying...")
    return [TextContent(type="text", text=json.dumps({"error": last_error}))]

async def main():
    async with stdio_server() as (read_stream, write_stream):
        await server.run(read_stream, write_stream, server.create_initialization_options())

if __name__ == "__main__":
    asyncio.run(main())