lambda_function.py

"""AWS GuardDuty Lambda Function Handler.

Processes GuardDuty findings from EventBridge, runs AI triage via
Bedrock AgentCore Runtime, and sends enriched notifications via
SNS, Email, and Google Chat.

Environment Variables:
    SNS_TOPIC_ARN: ARN of SNS topic for notifications
    GOOGLE_CHAT_WEBHOOK: Google Chat webhook URL
    ENABLE_SNS: Enable SNS notifications (true/false)
    ENABLE_EMAIL: Enable email notifications (true/false)
    ENABLE_CHAT: Enable Google Chat notifications (true/false)
    AGENTCORE_AGENT_ARN: Full ARN of the AgentCore triage agent runtime
    AGENTCORE_FORENSICS_ARN: Full ARN of the AgentCore forensics agent runtime
    AWS_REGION_NAME: AWS region for AgentCore client (default: ap-south-1)
"""

import json
import os
import logging
import uuid
from typing import Dict, Any, Optional
import boto3
from botocore.config import Config
import requests
from AWSSession import get_aws_session
from Notification import send_email

# Configure logging
logger = logging.getLogger()
logger.setLevel(logging.INFO)


def lambda_handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]:
    """
    Process GuardDuty findings and send multi-channel notifications.
    
    Args:
        event: EventBridge event containing GuardDuty finding
        context: Lambda context object
        
    Returns:
        Dict containing response status and details
    """
    logger.info("Processing GuardDuty finding notification")
    logger.info(f"Received event: {json.dumps(event)}")
    
    try:
        # Load configuration
        logger.info("Loading configuration from input.json")
        config = _load_configuration()
        logger.info("Configuration loaded successfully")
        
        # Extract finding details
        logger.info("Extracting finding data from event")
        finding_data = _extract_finding_data(event)
        logger.info(f"Finding data extracted: {json.dumps(finding_data)}")

        # Run AI triage via AgentCore Runtime
        logger.info("Invoking AgentCore triage agent")
        triage_result = _invoke_triage_agent(event.get('detail', {}))
        logger.info(f"Triage result: {json.dumps(triage_result)}")

        # Conditionally run deep forensics investigation
        # Triggers only when: requires_escalation=True AND calculated_risk>=High AND base_severity>=Low
        forensics_result = None
        if _should_run_forensics(triage_result):
            logger.info("Triage conditions met — invoking Deep Forensics Agent")
            forensics_result = _invoke_forensics_agent(event.get('detail', {}), triage_result)
            logger.info(f"Forensics result: {json.dumps(forensics_result)}")
        else:
            logger.info("Forensics conditions not met — skipping deep investigation")

        # Trigger AI Voice SOC Escalation when requires_escalation=True
        escalation_execution_arn = None
        if triage_result and triage_result.get('risk_assessment', {}).get('requires_escalation'):
            logger.info("requires_escalation=True — triggering voice SOC escalation workflow")
            escalation_execution_arn = _trigger_voice_escalation(
                finding_id=finding_data['finding_id'],
                finding_detail=event.get('detail', {}),
                triage_result=triage_result,
            )
            logger.info(f"Escalation workflow started: {escalation_execution_arn}")

        # Send notifications based on configuration
        logger.info("Starting notification delivery")
        notification_results = _send_notifications(config, finding_data, triage_result, forensics_result)
        logger.info(f"Notification results: {json.dumps(notification_results)}")
        
        logger.info("GuardDuty notification processing completed successfully")
        
        return {
            'statusCode': 200,
            'body': json.dumps({
                'message': 'Notifications processed successfully',
                'results': notification_results,
                'finding_id': finding_data['finding_id'],
                'triage': triage_result,
                'forensics': forensics_result,
                'escalation_execution_arn': escalation_execution_arn,
            })
        }
        
    except Exception as e:
        logger.error(f"Failed to process GuardDuty notification: {str(e)}")
        return {
            'statusCode': 500,
            'body': json.dumps({
                'error': 'Notification processing failed',
                'message': str(e)
            })
        }


def _invoke_triage_agent(finding_detail: Dict[str, Any]) -> Optional[Dict[str, Any]]:
    """Invoke the AgentCore triage agent and return the parsed risk assessment.

    Falls back gracefully to None if the call fails so notifications still fire.
    """
    agent_arn = os.environ.get('AGENTCORE_AGENT_ARN')
    region = os.environ.get('AWS_REGION_NAME', 'ap-south-1')

    if not agent_arn:
        logger.warning("AGENTCORE_AGENT_ARN not set — skipping triage")
        return None

    try:
        client = boto3.client('bedrock-agentcore', region_name=region)

        # runtimeSessionId must be 33+ characters
        session_id = str(uuid.uuid4()) + str(uuid.uuid4())[:4]

        prompt = (
            "Analyze the following GuardDuty finding and return the JSON risk assessment:\n\n"
            + json.dumps(finding_detail, indent=2)
        )

        response = client.invoke_agent_runtime(
            agentRuntimeArn=agent_arn,
            runtimeSessionId=session_id,
            payload=json.dumps({"prompt": prompt}),
        )

        raw_output = response['response'].read().decode('utf-8')
        if not raw_output:
            logger.warning("AgentCore returned empty response")
            return None

        result = json.loads(raw_output)

        # Agent returns {"triage": <parsed dict>, "raw": "<json string>"}
        triage = result.get('triage', result)

        # Guard: if triage is still a Strands message object instead of the assessment dict
        # (can happen if agent code wasn't redeployed), extract and parse the text
        if isinstance(triage, dict) and 'content' in triage:
            content = triage.get('content', [])
            text = content[0].get('text', '') if content else ''
            clean = text.strip()
            if clean.startswith('```'):
                clean = clean.split('```')[1]
                if clean.startswith('json'):
                    clean = clean[4:]
                clean = clean.strip()
            triage = json.loads(clean)

        logger.info(
            "Triage completed: risk=%s escalate=%s confidence=%s",
            triage.get('risk_assessment', {}).get('calculated_risk'),
            triage.get('risk_assessment', {}).get('requires_escalation'),
            triage.get('confidence_score'),
        )
        return triage

    except Exception as exc:
        logger.error(f"AgentCore triage invocation failed: {exc}")
        return None


_RISK_LEVELS = {"Low": 0, "Medium": 1, "High": 2, "Critical": 3}


def _should_run_forensics(triage_result: Optional[Dict[str, Any]]) -> bool:
    """Return True only when ALL three conditions are met:
      - requires_escalation == True
      - calculated_risk >= High
      - base_severity >= Low
    """
    if not triage_result:
        return False
    risk = triage_result.get("risk_assessment", {})
    requires_escalation = risk.get("requires_escalation", False)
    calculated_risk = _RISK_LEVELS.get(risk.get("calculated_risk", "Low"), 0)
    base_severity = _RISK_LEVELS.get(risk.get("base_severity", "Low"), 0)
    return (
        requires_escalation is True
        and calculated_risk >= _RISK_LEVELS["High"]
        and base_severity >= _RISK_LEVELS["Low"]
    )


def _invoke_forensics_agent(
    finding_detail: Dict[str, Any],
    triage_result: Dict[str, Any],
) -> Optional[Dict[str, Any]]:
    """Invoke the AgentCore forensics agent for deep investigation.

    Falls back gracefully to None if the call fails so notifications still fire.
    """
    agent_arn = os.environ.get("AGENTCORE_FORENSICS_ARN")
    region = os.environ.get("AWS_REGION_NAME", "ap-south-1")

    if not agent_arn:
        logger.warning("AGENTCORE_FORENSICS_ARN not set — skipping forensics")
        return None

    try:
        client = boto3.client(
            "bedrock-agentcore",
            region_name=region,
            config=Config(
                connect_timeout=10,
                read_timeout=300,
                retries={"max_attempts": 0},
            ),
        )
        session_id = str(uuid.uuid4()) + str(uuid.uuid4())[:4]

        response = client.invoke_agent_runtime(
            agentRuntimeArn=agent_arn,
            runtimeSessionId=session_id,
            payload=json.dumps({"finding": finding_detail, "triage": triage_result}),
        )

        raw_output = response["response"].read().decode("utf-8")
        if not raw_output:
            logger.warning("Forensics AgentCore returned empty response")
            return None

        result = json.loads(raw_output)
        forensics = result.get("forensics", result)

        if isinstance(forensics, dict) and "content" in forensics:
            content = forensics.get("content", [])
            text = content[0].get("text", "") if content else ""
            clean = text.strip()
            if clean.startswith("```"):
                clean = clean.split("```")[1]
                if clean.startswith("json"):
                    clean = clean[4:]
                clean = clean.strip()
            forensics = json.loads(clean)

        verdict = forensics.get("overall_verdict", {})
        logger.info(
            "Forensics completed: compromise=%s stage=%s confidence=%s",
            verdict.get("confirmed_compromise"),
            verdict.get("attack_stage"),
            verdict.get("confidence"),
        )
        return forensics

    except Exception as exc:
        logger.error(f"Forensics agent invocation failed: {exc}")
        return None


def _trigger_voice_escalation(
    finding_id: str,
    finding_detail: Dict[str, Any],
    triage_result: Dict[str, Any],
) -> Optional[str]:
    """Start the AI Voice SOC Escalation Step Functions state machine.

    Returns the execution ARN on success, or None if the env var is not set
    or the call fails (so existing notifications still fire).
    """
    sfn_arn = os.environ.get("ESCALATION_STATE_MACHINE_ARN")
    if not sfn_arn:
        logger.warning("ESCALATION_STATE_MACHINE_ARN not set — skipping voice escalation")
        return None

    try:
        sfn = boto3.client("stepfunctions")
        payload = json.dumps({
            "finding_id": finding_id,
            "finding_detail": finding_detail,
            "triage": triage_result,
        })
        response = sfn.start_execution(
            stateMachineArn=sfn_arn,
            name=f"escalation-{finding_id[:36]}",
            input=payload,
        )
        return response["executionArn"]
    except Exception as exc:
        logger.error(f"Failed to start voice escalation workflow: {exc}")
        return None


def _load_configuration() -> Dict[str, Any]:
    try:
        with open('input.json', 'r') as f:
            config = json.load(f)
            logger.info("Configuration file read successfully")
            return config
    except Exception as e:
        logger.error(f"Failed to load configuration: {str(e)}")
        raise


def _extract_finding_data(event: Dict[str, Any]) -> Dict[str, Any]:
    """Extract relevant data from GuardDuty finding event."""
    detail = event['detail']
    logger.info(f"Extracting data from finding ID: {detail.get('id', 'unknown')}")
    
    return {
        'severity': detail['severity'],
        'title': detail['title'],
        'description': detail['description'],
        'finding_type': detail['type'],
        'resource': detail['resource']['resourceType'],
        'account': event['account'],
        'region': event['region'],
        'finding_id': detail['id'],
        'created_at': detail.get('createdAt', ''),
        'updated_at': detail.get('updatedAt', '')
    }


def _send_notifications(config: Dict[str, Any], finding_data: Dict[str, Any], triage_result: Optional[Dict[str, Any]], forensics_result: Optional[Dict[str, Any]] = None) -> Dict[str, str]:
    """Send notifications via configured channels."""
    results = {}
    
    # SNS Notification
    if _is_enabled('ENABLE_SNS') and os.environ.get('SNS_TOPIC_ARN'):
        logger.info("SNS notification enabled, attempting to send")
        try:
            _send_sns_notification(config, finding_data, triage_result, forensics_result)
            results['sns'] = 'Success'
            logger.info("SNS notification completed successfully")
        except Exception as e:
            logger.error(f"SNS notification failed: {str(e)}")
            results['sns'] = f'Failed: {str(e)}'
    
    # Email Notification
    if _is_enabled('ENABLE_EMAIL'):
        logger.info("Email notification enabled, attempting to send")
        try:
            _send_email_notification(config, finding_data, triage_result, forensics_result)
            results['email'] = 'Success'
            logger.info("Email notification completed successfully")
        except Exception as e:
            logger.error(f"Email notification failed: {str(e)}")
            results['email'] = f'Failed: {str(e)}'
    
    # Google Chat Notification
    if _is_enabled('ENABLE_CHAT') and os.environ.get('GOOGLE_CHAT_WEBHOOK'):
        logger.info("Google Chat notification enabled, attempting to send")
        try:
            _send_chat_notification(finding_data, triage_result, forensics_result)
            results['chat'] = 'Success'
            logger.info("Google Chat notification completed successfully")
        except Exception as e:
            logger.error(f"Chat notification failed: {str(e)}")
            results['chat'] = f'Failed: {str(e)}'
    
    return results


def _send_sns_notification(config: Dict[str, Any], finding_data: Dict[str, Any], triage_result: Optional[Dict[str, Any]], forensics_result: Optional[Dict[str, Any]] = None) -> None:
    """Send SNS notification."""
    logger.info("Creating AWS session for SNS")
    session = get_aws_session(config['awsCredentials'])
    sns_client = session.client('sns')
    
    logger.info("Formatting SNS message")
    message = _format_sns_message(finding_data, triage_result, forensics_result)
    subject = f"GuardDuty Alert [{finding_data['severity']}] | {finding_data['title']}"
    subject = subject[:100]
    logger.info(f"SNS subject: {subject}")
    
    logger.info(f"Publishing to SNS topic: {os.environ['SNS_TOPIC_ARN']}")
    sns_client.publish(
        TopicArn=os.environ['SNS_TOPIC_ARN'],
        Subject=subject,
        Message=message
    )
    
    logger.info("SNS notification sent successfully")


def _send_email_notification(config: Dict[str, Any], finding_data: Dict[str, Any], triage_result: Optional[Dict[str, Any]], forensics_result: Optional[Dict[str, Any]] = None) -> None:
    """Send email notification."""
    logger.info("Formatting email content")
    email_content = _format_email_content(finding_data, triage_result, forensics_result)
    
    email_details = config['emailNotification'].copy()
    email_details['email_subject'] = f"GuardDuty Alert | {finding_data['title']}"
    logger.info(f"Email subject: {email_details['email_subject']}")
    logger.info(f"Email recipients: {email_details.get('email_to', [])}")
    
    logger.info("Sending email via SMTP")
    send_email(config['smtpCredentials'], email_details, email_content)
    logger.info("Email notification sent successfully")


def _send_chat_notification(finding_data: Dict[str, Any], triage_result: Optional[Dict[str, Any]], forensics_result: Optional[Dict[str, Any]] = None) -> None:
    """Send Google Chat notification."""
    logger.info("Formatting Google Chat message")
    message = _format_chat_message(finding_data, triage_result, forensics_result)
    
    webhook_url = os.environ['GOOGLE_CHAT_WEBHOOK']
    logger.info(f"Posting to Google Chat webhook: {webhook_url[:50]}...")
    response = requests.post(
        webhook_url,
        json=message,
        timeout=10
    )
    
    response.raise_for_status()
    logger.info(f"Google Chat response status: {response.status_code}")
    logger.info("Google Chat notification sent successfully")


def _format_sns_message(finding_data: Dict[str, Any], triage_result: Optional[Dict[str, Any]], forensics_result: Optional[Dict[str, Any]] = None) -> str:
    """Format concise message for SNS notification."""
    base = f"""GuardDuty Security Alert

Title      : {finding_data['title']}
Description: {finding_data['description']}
Type       : {finding_data['finding_type']}
Severity   : {finding_data['severity']}
Resource   : {finding_data['resource']}
Account    : {finding_data['account']}
Region     : {finding_data['region']}
Finding ID : {finding_data['finding_id']}"""

    if triage_result:
        risk = triage_result.get('risk_assessment', {})
        threat = triage_result.get('threat_analysis', {})
        escalation_flag = "⚠️ ESCALATION REQUIRED" if risk.get('requires_escalation') else "No escalation needed"
        top_containment = triage_result.get('immediate_containment', [])[:2]
        base += f"""

── AI Triage ─────────────────────────────────────
Risk      : {risk.get('calculated_risk', 'N/A')} (base: {risk.get('base_severity', 'N/A')})
Escalation: {escalation_flag}
Threats   : {', '.join(threat.get('category', []))}
Confidence: {triage_result.get('confidence_score', 'N/A')}
Summary   : {triage_result.get('finding_summary', 'N/A')}

Top Containment Steps:
{chr(10).join(f'  • {s}' for s in top_containment)}
──────────────────────────────────────────────────

Action Required: Investigate in AWS GuardDuty console immediately."""

    if forensics_result:
        verdict = forensics_result.get('overall_verdict', {})
        lateral = forensics_result.get('lateral_movement_assessment', {})
        persistence = forensics_result.get('persistence_assessment', {})
        top_actions = forensics_result.get('prioritized_actions', [])[:3]
        actions_str = chr(10).join(
            f"  {a.get('priority', '?')}. {a.get('action', '')}" for a in top_actions
        )
        base += f"""

── 🔬 Deep Forensics ─────────────────────────────
Summary   : {forensics_result.get('forensics_summary', 'N/A')}
Compromise: {'🔴 CONFIRMED' if verdict.get('confirmed_compromise') else '🟡 Unconfirmed'}
Stage     : {verdict.get('attack_stage', 'N/A')}
Confidence: {verdict.get('confidence', 'N/A')}
Lateral   : {'🔴 DETECTED' if lateral.get('detected') else '✅ None'}
Persistence: {'🔴 DETECTED' if persistence.get('detected') else '✅ None'}

Top Priority Actions:
{actions_str}
──────────────────────────────────────────────────"""

    return base


def _format_email_content(finding_data: Dict[str, Any], triage_result: Optional[Dict[str, Any]], forensics_result: Optional[Dict[str, Any]] = None) -> str:
    """Format full HTML email with all triage keys."""
    severity_color = _get_severity_color(finding_data['severity'])

    triage_section = ""
    if triage_result:
        risk   = triage_result.get('risk_assessment', {})
        threat = triage_result.get('threat_analysis', {})
        res    = triage_result.get('resource', {})
        auto   = triage_result.get('automation_feasibility', {})

        calculated_risk = risk.get('calculated_risk', 'N/A')
        risk_color = {"Low": "#28a745", "Medium": "#fd7e14", "High": "#dc3545", "Critical": "#6f0000"}.get(calculated_risk, "#6c757d")

        escalation_badge = (
            '<span style="background:#dc3545;color:#fff;padding:2px 8px;border-radius:4px;font-size:12px;">⚠️ ESCALATION REQUIRED</span>'
            if risk.get('requires_escalation')
            else '<span style="background:#28a745;color:#fff;padding:2px 8px;border-radius:4px;font-size:12px;">✅ No Escalation</span>'
        )
        auto_badge = (
            '<span style="background:#28a745;color:#fff;padding:2px 8px;border-radius:4px;font-size:12px;">✅ Yes</span>'
            if auto.get('can_auto_remediate')
            else '<span style="background:#6c757d;color:#fff;padding:2px 8px;border-radius:4px;font-size:12px;">Manual Required</span>'
        )

        def ul(items): return "<ul style='margin:4px 0;padding-left:18px;'>" + "".join(f"<li>{i}</li>" for i in items) + "</ul>" if items else "N/A"

        # Threat flags row
        flags = []
        if threat.get('persistence_detected'):    flags.append('🔴 Persistence')
        if threat.get('lateral_movement_risk'):   flags.append('🔴 Lateral Movement')
        if threat.get('data_exfiltration_risk'):  flags.append('🔴 Data Exfiltration')
        if res.get('public_exposure'):            flags.append('🔴 Public Exposure')
        flags_html = ' &nbsp;|&nbsp; '.join(flags) if flags else '✅ None detected'

        # Environment tags
        env_tags = res.get('environment_tags', {})
        env_tags_html = ', '.join(f"{k}={v}" for k, v in env_tags.items()) if isinstance(env_tags, dict) and env_tags else str(env_tags) if env_tags else 'N/A'

        td_label = "padding:8px;font-weight:bold;background:#f5f5f5;white-space:nowrap;vertical-align:top;width:160px;"
        td_value = "padding:8px;vertical-align:top;"

        triage_section = f"""
            <h3 style="color:#333;margin-top:30px;border-bottom:2px solid #dee2e6;padding-bottom:8px;">🤖 AI Triage Assessment</h3>

            <table style="border-collapse:collapse;width:100%;margin-top:10px;">
                <tr><td style="{td_label}">Summary</td>
                    <td style="{td_value}">{triage_result.get('finding_summary', 'N/A')}</td></tr>
                <tr><td style="{td_label}">Calculated Risk</td>
                    <td style="{td_value};color:{risk_color};font-weight:bold;">{calculated_risk}</td></tr>
                <tr><td style="{td_label}">Base Severity</td>
                    <td style="{td_value}">{risk.get('base_severity', 'N/A')}</td></tr>
                <tr><td style="{td_label}">Escalation</td>
                    <td style="{td_value}">{escalation_badge}</td></tr>
                <tr><td style="{td_label}">Blast Radius</td>
                    <td style="{td_value}">{risk.get('blast_radius_estimate', 'N/A')}</td></tr>
                <tr><td style="{td_label}">Reasoning</td>
                    <td style="{td_value}">{risk.get('reasoning', 'N/A')}</td></tr>
                <tr><td style="{td_label}">Confidence</td>
                    <td style="{td_value}">{triage_result.get('confidence_score', 'N/A')}</td></tr>
            </table>

            <h4 style="color:#495057;margin-top:20px;">🎯 Threat Analysis</h4>
            <table style="border-collapse:collapse;width:100%;">
                <tr><td style="{td_label}">Categories</td>
                    <td style="{td_value}">{', '.join(threat.get('category', []))}</td></tr>
                <tr><td style="{td_label}">Attack Vector</td>
                    <td style="{td_value}">{threat.get('attack_vector', 'N/A')}</td></tr>
                <tr><td style="{td_label}">Privilege Level</td>
                    <td style="{td_value}">{threat.get('privilege_level', 'N/A')}</td></tr>
                <tr><td style="{td_label}">Risk Flags</td>
                    <td style="{td_value}">{flags_html}</td></tr>
            </table>

            <h4 style="color:#495057;margin-top:20px;">🖥️ Affected Resource</h4>
            <table style="border-collapse:collapse;width:100%;">
                <tr><td style="{td_label}">Type / ID</td>
                    <td style="{td_value}">{res.get('type', 'N/A')} / {res.get('id', 'N/A')}</td></tr>
                <tr><td style="{td_label}">Region / Account</td>
                    <td style="{td_value}">{res.get('region', 'N/A')} / {res.get('account_id', 'N/A')}</td></tr>
                <tr><td style="{td_label}">Public Exposure</td>
                    <td style="{td_value}">{'🔴 Yes' if res.get('public_exposure') else '✅ No'}</td></tr>
                <tr><td style="{td_label}">Environment Tags</td>
                    <td style="{td_value}">{env_tags_html}</td></tr>
            </table>

            <h4 style="color:#495057;margin-top:20px;">💼 Business Impact</h4>
            <p style="margin:4px 0;padding:8px;background:#fff3cd;border-radius:4px;">{triage_result.get('business_impact', 'N/A')}</p>

            <h4 style="color:#c0392b;margin-top:20px;">🚨 Immediate Containment</h4>
            {ul(triage_result.get('immediate_containment', []))}

            <h4 style="color:#2980b9;margin-top:15px;">🔍 Investigation Steps</h4>
            {ul(triage_result.get('investigation_steps', []))}

            <h4 style="color:#8e44ad;margin-top:15px;">🔧 Remediation Actions</h4>
            {ul(triage_result.get('remediation_actions', []))}

            <h4 style="color:#16a085;margin-top:15px;">🛡️ Long-Term Prevention</h4>
            {ul(triage_result.get('long_term_prevention', []))}

            <h4 style="color:#495057;margin-top:15px;">⚙️ Automation Feasibility</h4>
            <table style="border-collapse:collapse;width:100%;">
                <tr><td style="{td_label}">Can Auto-Remediate</td>
                    <td style="{td_value}">{auto_badge}</td></tr>
                <tr><td style="{td_label}">Remediation Risk</td>
                    <td style="{td_value}">{auto.get('risk_of_auto_remediation', 'N/A')}</td></tr>
                <tr><td style="{td_label}">Recommended</td>
                    <td style="{td_value}">{ul(auto.get('recommended_automation', []))}</td></tr>
            </table>

            <h4 style="color:#495057;margin-top:15px;">🗺️ MITRE ATT&CK Mapping</h4>
            {ul(triage_result.get('mitre_attack_mapping', []))}"""

    # ── Forensics section ──────────────────────────────────────────────────────
    forensics_section = ""
    if forensics_result:
        verdict = forensics_result.get('overall_verdict', {})
        lateral = forensics_result.get('lateral_movement_assessment', {})
        persistence = forensics_result.get('persistence_assessment', {})
        container = forensics_result.get('container_security_posture', {})
        iam_blast = forensics_result.get('iam_blast_radius', {})
        network = forensics_result.get('network_analysis', {})
        ct = forensics_result.get('cloudtrail_findings', {})
        lineage = forensics_result.get('process_lineage_analysis', {})

        compromise_badge = (
            '<span style="background:#6f0000;color:#fff;padding:2px 8px;border-radius:4px;">🔴 CONFIRMED COMPROMISE</span>'
            if verdict.get('confirmed_compromise')
            else '<span style="background:#fd7e14;color:#fff;padding:2px 8px;border-radius:4px;">🟡 Unconfirmed</span>'
        )
        escape_color = {"None": "#28a745", "Possible": "#fd7e14", "Likely": "#dc3545", "Confirmed": "#6f0000"}.get(
            container.get('escape_feasibility', 'None'), "#6c757d"
        )

        def ul_f(items): return "<ul style='margin:4px 0;padding-left:18px;'>" + "".join(f"<li>{i}</li>" for i in items) + "</ul>" if items else "<em>None</em>"

        priority_rows = "".join(
            f"<tr><td style='padding:4px 8px;font-weight:bold;'>#{a.get('priority','?')}</td>"
            f"<td style='padding:4px 8px;'>{a.get('action','')}</td>"
            f"<td style='padding:4px 8px;color:#6c757d;font-size:12px;'>{a.get('rationale','')}</td></tr>"
            for a in forensics_result.get('prioritized_actions', [])
        )

        forensics_section = f"""
            <h3 style="color:#6f0000;margin-top:40px;border-bottom:3px solid #6f0000;padding-bottom:8px;">🔬 Deep Forensics Investigation</h3>

            <table style="border-collapse:collapse;width:100%;margin-top:10px;">
                <tr><td style="{td_label}">Summary</td>
                    <td style="{td_value}">{forensics_result.get('forensics_summary', 'N/A')}</td></tr>
                <tr><td style="{td_label}">Compromise</td>
                    <td style="{td_value}">{compromise_badge}</td></tr>
                <tr><td style="{td_label}">Attack Stage</td>
                    <td style="{td_value};font-weight:bold;">{verdict.get('attack_stage', 'N/A')}</td></tr>
                <tr><td style="{td_label}">Attacker Goal</td>
                    <td style="{td_value}">{verdict.get('attacker_objectives', 'N/A')}</td></tr>
                <tr><td style="{td_label}">Dwell Time</td>
                    <td style="{td_value}">{verdict.get('dwell_time_estimate', 'N/A')}</td></tr>
                <tr><td style="{td_label}">Confidence</td>
                    <td style="{td_value}">{verdict.get('confidence', 'N/A')}</td></tr>
            </table>

            <h4 style="color:#495057;margin-top:20px;">🧬 Process Lineage</h4>
            <table style="border-collapse:collapse;width:100%;">
                <tr><td style="{td_label}">Chain</td>
                    <td style="{td_value};font-family:monospace;font-size:12px;">{lineage.get('chain', 'N/A')}</td></tr>
                <tr><td style="{td_label}">Attack Pattern</td>
                    <td style="{td_value}">{lineage.get('attack_pattern', 'None')}</td></tr>
                <tr><td style="{td_label}">Root Execution</td>
                    <td style="{td_value}">{'🔴 Yes' if lineage.get('root_execution') else '✅ No'}</td></tr>
                <tr><td style="{td_label}">Escape Indicators</td>
                    <td style="{td_value}">{'🔴 Yes' if lineage.get('container_escape_indicators') else '✅ No'}</td></tr>
                <tr><td style="{td_label}">Patterns</td>
                    <td style="{td_value}">{ul_f(lineage.get('suspicious_patterns', []))}</td></tr>
            </table>

            <h4 style="color:#495057;margin-top:20px;">🐳 Container Security Posture</h4>
            <table style="border-collapse:collapse;width:100%;">
                <tr><td style="{td_label}">Privileged Mode</td>
                    <td style="{td_value}">{'🔴 Yes' if container.get('privileged_mode') else '✅ No'}</td></tr>
                <tr><td style="{td_label}">Host Network</td>
                    <td style="{td_value}">{'🔴 Yes' if container.get('host_network') else '✅ No'}</td></tr>
                <tr><td style="{td_label}">Host PID</td>
                    <td style="{td_value}">{'🔴 Yes' if container.get('host_pid') else '✅ No'}</td></tr>
                <tr><td style="{td_label}">Dangerous Caps</td>
                    <td style="{td_value}">{', '.join(container.get('dangerous_capabilities', [])) or '✅ None'}</td></tr>
                <tr><td style="{td_label}">Escape Feasibility</td>
                    <td style="{td_value};color:{escape_color};font-weight:bold;">{container.get('escape_feasibility', 'N/A')}</td></tr>
                <tr><td style="{td_label}">Escape Reasoning</td>
                    <td style="{td_value}">{container.get('escape_reasoning', 'N/A')}</td></tr>
            </table>

            <h4 style="color:#495057;margin-top:20px;">🔑 IAM Blast Radius</h4>
            <table style="border-collapse:collapse;width:100%;">
                <tr><td style="{td_label}">Instance Role</td>
                    <td style="{td_value};font-family:monospace;">{iam_blast.get('instance_role', 'N/A')}</td></tr>
                <tr><td style="{td_label}">Task Role</td>
                    <td style="{td_value};font-family:monospace;">{iam_blast.get('task_role', 'N/A')}</td></tr>
                <tr><td style="{td_label}">Blast Scope</td>
                    <td style="{td_value};font-weight:bold;">{iam_blast.get('blast_radius_scope', 'N/A')}</td></tr>
                <tr><td style="{td_label}">Can Escalate</td>
                    <td style="{td_value}">{'🔴 Yes' if iam_blast.get('can_escalate_privileges') else '✅ No'}</td></tr>
                <tr><td style="{td_label}">Can Access Secrets</td>
                    <td style="{td_value}">{'🔴 Yes' if iam_blast.get('can_access_secrets') else '✅ No'}</td></tr>
                <tr><td style="{td_label}">Can Move Laterally</td>
                    <td style="{td_value}">{'🔴 Yes' if iam_blast.get('can_move_laterally') else '✅ No'}</td></tr>
                <tr><td style="{td_label}">Overpermissive Actions</td>
                    <td style="{td_value}">{ul_f(iam_blast.get('overly_permissive_actions', []))}</td></tr>
            </table>

            <h4 style="color:#495057;margin-top:20px;">🌐 Network Analysis</h4>
            <table style="border-collapse:collapse;width:100%;">
                <tr><td style="{td_label}">SG Issues</td>
                    <td style="{td_value}">{ul_f(network.get('security_group_issues', []))}</td></tr>
                <tr><td style="{td_label}">Suspicious IPs</td>
                    <td style="{td_value}">{', '.join(network.get('suspicious_outbound_ips', [])) or '✅ None'}</td></tr>
                <tr><td style="{td_label}">C2 Beaconing</td>
                    <td style="{td_value}">{'🔴 Detected' if network.get('c2_beaconing_indicators') else '✅ None'}</td></tr>
                <tr><td style="{td_label}">Data Exfiltration</td>
                    <td style="{td_value}">{'🔴 Detected' if network.get('data_exfiltration_indicators') else '✅ None'}</td></tr>
            </table>

            <h4 style="color:#495057;margin-top:20px;">📋 CloudTrail Findings</h4>
            <table style="border-collapse:collapse;width:100%;">
                <tr><td style="{td_label}">Recon API Calls</td>
                    <td style="{td_value}">{ul_f(ct.get('recon_api_calls', []))}</td></tr>
                <tr><td style="{td_label}">Privilege Escalation</td>
                    <td style="{td_value}">{ul_f(ct.get('privilege_escalation_attempts', []))}</td></tr>
                <tr><td style="{td_label}">Data Access</td>
                    <td style="{td_value}">{ul_f(ct.get('data_access_events', []))}</td></tr>
            </table>

            <h4 style="color:#c0392b;margin-top:20px;">🚨 Lateral Movement</h4>
            <p>{'🔴 <strong>DETECTED</strong>' if lateral.get('detected') else '✅ None detected'}</p>
            {ul_f(lateral.get('evidence', []))}

            <h4 style="color:#8e44ad;margin-top:15px;">🔒 Persistence</h4>
            <p>{'🔴 <strong>DETECTED</strong>' if persistence.get('detected') else '✅ None detected'}</p>
            {ul_f(persistence.get('techniques', []))}

            <h4 style="color:#2980b9;margin-top:15px;">⚡ Prioritized Actions</h4>
            <table style="border-collapse:collapse;width:100%;border:1px solid #dee2e6;">
                <tr style="background:#f5f5f5;"><th style="padding:6px 8px;text-align:left;">#</th>
                    <th style="padding:6px 8px;text-align:left;">Action</th>
                    <th style="padding:6px 8px;text-align:left;">Rationale</th></tr>
                {priority_rows}
            </table>

            <h4 style="color:#495057;margin-top:15px;">🗺️ MITRE ATT&CK (Forensics)</h4>
            {ul_f(forensics_result.get('mitre_attack_mapping', []))}

            <h4 style="color:#6c757d;margin-top:15px;">⚠️ Data Gaps</h4>
            {ul_f(forensics_result.get('data_gaps', []))}"""

    return f"""
    <html>
    <body style="font-family:Arial,sans-serif;margin:20px;color:#212529;">
        <div style="border-left:4px solid {severity_color};padding-left:20px;">
            <h2 style="color:{severity_color};">🚨 GuardDuty Security Alert</h2>

            <table style="border-collapse:collapse;width:100%;margin-top:20px;">
                <tr><td style="padding:8px;font-weight:bold;background:#f5f5f5;">Title</td>
                    <td style="padding:8px;">{finding_data['title']}</td></tr>
                <tr><td style="padding:8px;font-weight:bold;background:#f5f5f5;">Description</td>
                    <td style="padding:8px;">{finding_data['description']}</td></tr>
                <tr><td style="padding:8px;font-weight:bold;background:#f5f5f5;">Type</td>
                    <td style="padding:8px;">{finding_data['finding_type']}</td></tr>
                <tr><td style="padding:8px;font-weight:bold;background:#f5f5f5;">Severity</td>
                    <td style="padding:8px;color:{severity_color};font-weight:bold;">{finding_data['severity']}</td></tr>
                <tr><td style="padding:8px;font-weight:bold;background:#f5f5f5;">Resource</td>
                    <td style="padding:8px;">{finding_data['resource']}</td></tr>
                <tr><td style="padding:8px;font-weight:bold;background:#f5f5f5;">Account</td>
                    <td style="padding:8px;">{finding_data['account']}</td></tr>
                <tr><td style="padding:8px;font-weight:bold;background:#f5f5f5;">Region</td>
                    <td style="padding:8px;">{finding_data['region']}</td></tr>
                <tr><td style="padding:8px;font-weight:bold;background:#f5f5f5;">Finding ID</td>
                    <td style="padding:8px;font-family:monospace;">{finding_data['finding_id']}</td></tr>
            </table>

            {triage_section}

            {forensics_section}

            <div style="margin-top:20px;padding:15px;background:#fff3cd;border:1px solid #ffeaa7;border-radius:4px;">
                <strong>Action Required:</strong> Investigate this finding in the AWS GuardDuty console immediately.
            </div>
        </div>
    </body>
    </html>
    """


def _format_chat_message(finding_data: Dict[str, Any], triage_result: Optional[Dict[str, Any]], forensics_result: Optional[Dict[str, Any]] = None) -> Dict[str, str]:
    """Format concise message for Google Chat notification."""
    triage_block = ""
    if triage_result:
        risk   = triage_result.get('risk_assessment', {})
        threat = triage_result.get('threat_analysis', {})
        escalation = "⚠️ *ESCALATION REQUIRED*" if risk.get('requires_escalation') else "✅ No escalation"

        flags = []
        if threat.get('persistence_detected'):   flags.append('Persistence')
        if threat.get('lateral_movement_risk'):  flags.append('Lateral Movement')
        if threat.get('data_exfiltration_risk'): flags.append('Data Exfiltration')
        flags_str = ', '.join(flags) if flags else 'None'

        top_containment = '\n'.join(
            f'  • {s}' for s in triage_result.get('immediate_containment', [])[:2]
        )

        triage_block = f"""

🤖 *AI Triage*
*Risk:* {risk.get('calculated_risk', 'N/A')} (base: {risk.get('base_severity', 'N/A')}) | *Confidence:* {triage_result.get('confidence_score', 'N/A')}
*Escalation:* {escalation}
*Threats:* {', '.join(threat.get('category', []))}
*Risk Flags:* {flags_str}
*Summary:* {triage_result.get('finding_summary', 'N/A')}

🚨 *Top Containment Steps:*
{top_containment}"""

    forensics_block = ""
    if forensics_result:
        verdict = forensics_result.get('overall_verdict', {})
        lateral = forensics_result.get('lateral_movement_assessment', {})
        persistence = forensics_result.get('persistence_assessment', {})
        container = forensics_result.get('container_security_posture', {})
        iam_blast = forensics_result.get('iam_blast_radius', {})
        top_actions = '\n'.join(
            f"  {a.get('priority','?')}. {a.get('action','')}"
            for a in forensics_result.get('prioritized_actions', [])[:3]
        )
        forensics_block = f"""

🔬 *Deep Forensics Investigation*
*Summary:* {forensics_result.get('forensics_summary', 'N/A')}
*Compromise:* {'🔴 CONFIRMED' if verdict.get('confirmed_compromise') else '🟡 Unconfirmed'} | *Stage:* {verdict.get('attack_stage', 'N/A')}
*Confidence:* {verdict.get('confidence', 'N/A')} | *Dwell Time:* {verdict.get('dwell_time_estimate', 'N/A')}
*Lateral Movement:* {'🔴 DETECTED' if lateral.get('detected') else '✅ None'} | *Persistence:* {'🔴 DETECTED' if persistence.get('detected') else '✅ None'}
*Container Escape:* {container.get('escape_feasibility', 'N/A')} | *IAM Blast Scope:* {iam_blast.get('blast_radius_scope', 'N/A')}

⚡ *Priority Actions:*
{top_actions}"""

    return {
        "text": f"""🚨 *GuardDuty Security Alert*

*Title:* {finding_data['title']}
*Description:* {finding_data['description']}
*Type:* {finding_data['finding_type']}
*Severity:* {finding_data['severity']}
*Resource:* {finding_data['resource']}
*Account:* {finding_data['account']}
*Region:* {finding_data['region']}
*Finding ID:* `{finding_data['finding_id']}`{triage_block}{forensics_block}

⚠️ *Action Required:* Investigate in AWS GuardDuty console immediately."""
    }


def _get_severity_color(severity: float) -> str:
    """Get color code based on severity level."""
    if severity >= 7.0:
        return "#dc3545"  # Red for high severity
    elif severity >= 4.0:
        return "#fd7e14"  # Orange for medium severity
    else:
        return "#28a745"  # Green for low severity


def _is_enabled(env_var: str) -> bool:
    """Check if notification channel is enabled."""
    return os.environ.get(env_var, 'false').lower() == 'true'