Stop Writing Runbooks. Start Building Self-Healing Systems.

# remediation_router.py
import boto3
import json
import os
from datetime import datetime, timezone

ssm = boto3.client("ssm")
sns = boto3.client("sns")
dynamodb = boto3.resource("dynamodb")

# DynamoDB table tracks remediation history to detect repeated failures
history_table = dynamodb.Table(os.environ["REMEDIATION_HISTORY_TABLE"])
alert_topic_arn = os.environ["ALERT_SNS_TOPIC_ARN"]

# Map event sources to SSM Automation Document names and safe thresholds
REMEDIATION_CONFIG = {
"ec2-lifecycle-stuck": {
"ssm_document": "AutoHeal-EC2LifecycleStuck",
"max_auto_remediations_per_hour": 3, # Safety cap: don't loop forever
"escalate_after_attempts": 3,
},
"ecs-task-failure": {
"ssm_document": "AutoHeal-ECSTaskRestart",
"max_auto_remediations_per_hour": 5,
"escalate_after_attempts": 5,
},
"rds-storage-full": {
"ssm_document": "AutoHeal-RDSStorageExpand",
"max_auto_remediations_per_hour": 1, # Storage changes are expensive; be conservative
"escalate_after_attempts": 1,
},
}

def lambda_handler(event, context):
print(f"Received event: {json.dumps(event)}")

# Normalize the event into a structured remediation request
remediation_type, resource_id, metadata = classify_event(event)

if not remediation_type:
print("Event did not match any known remediation pattern. Ignoring.")
return {"status": "ignored"}

config = REMEDIATION_CONFIG[remediation_type]

# Check how many times we've already tried to auto-fix this resource
recent_attempts = get_recent_attempts(resource_id, remediation_type)
print(f"Recent attempts for {resource_id} ({remediation_type}): {recent_attempts}")

if recent_attempts >= config["escalate_after_attempts"]:
# We've tried too many times — escalate to a human instead
print(f"Escalation threshold reached for {resource_id}. Notifying team.")
notify_human(remediation_type, resource_id, recent_attempts, metadata)
return {"status": "escalated", "resource_id": resource_id}

# Check for active maintenance windows — don't auto-remediate during planned work
if is_maintenance_window_active(resource_id):
print(f"Maintenance window active for {resource_id}. Skipping auto-remediation.")
return {"status": "skipped_maintenance_window"}

# All checks passed — trigger SSM Automation
execution_id = trigger_ssm_automation(
document_name=config["ssm_document"],
resource_id=resource_id,
metadata=metadata,
)

# Record this attempt in DynamoDB for history tracking
record_attempt(resource_id, remediation_type, execution_id)

return {
"status": "remediation_triggered",
"resource_id": resource_id,
"ssm_execution_id": execution_id,
}

def classify_event(event):
"""
Inspect the EventBridge event and return (remediation_type, resource_id, metadata).
Returns (None, None, None) if no match is found.
"""
source = event.get("source", "")
detail = event.get("detail", {})

if source == "aws.autoscaling" and event.get("detail-type") == "EC2 Instance-launch Lifecycle Action":
instance_id = detail.get("EC2InstanceId")
lifecycle_hook = detail.get("LifecycleHookName")
asg_name = detail.get("AutoScalingGroupName")
token = detail.get("LifecycleActionToken")
return (
"ec2-lifecycle-stuck",
instance_id,
{"lifecycle_hook": lifecycle_hook, "asg_name": asg_name, "token": token},
)

if source == "aws.ecs" and event.get("detail-type") == "ECS Task State Change":
task_arn = detail.get("taskArn")
cluster_arn = detail.get("clusterArn")
task_def = detail.get("taskDefinitionArn")
stop_code = detail.get("stopCode")
return (
"ecs-task-failure",
task_arn,
{"cluster_arn": cluster_arn, "task_def": task_def, "stop_code": stop_code},
)

if source == "aws.rds" and detail.get("EventID") == "RDS-EVENT-0007":
db_instance_id = detail.get("SourceIdentifier")
return (
"rds-storage-full",
db_instance_id,
{"event_message": detail.get("Message", "")},
)

return None, None, None

def get_recent_attempts(resource_id, remediation_type, window_hours=1):
"""
Query DynamoDB for how many times we've attempted to remediate this
resource in the last `window_hours` hours.
"""
from boto3.dynamodb.conditions import Key, Attr
from datetime import timedelta

cutoff = datetime.now(timezone.utc) - timedelta(hours=window_hours)
cutoff_iso = cutoff.isoformat()

response = history_table.query(
KeyConditionExpression=Key("resource_id").eq(resource_id),
FilterExpression=Attr("remediation_type").eq(remediation_type)
& Attr("timestamp").gt(cutoff_iso),
)
return response["Count"]

def is_maintenance_window_active(resource_id):
"""
Check AWS Systems Manager Maintenance Windows to see if any active window
targets the given resource. This prevents auto-remediation during planned work.
"""
try:
response = ssm.describe_maintenance_window_executions(
Filters=[{"Key": "ExecutedAfter", "Values": [datetime.now(timezone.utc).isoformat()]}]
)
# A more complete implementation would check window targets against resource tags.
# For brevity, we return False (no maintenance window) if the API returns no active executions.
return len(response.get("WindowExecutions", [])) > 0
except Exception as e:
print(f"Could not check maintenance windows: {e}. Proceeding with remediation.")
return False

def trigger_ssm_automation(document_name, resource_id, metadata):
"""
Start an SSM Automation document execution with the resource context.
SSM Automation provides retries, approval gates, and a full audit trail.
"""
parameters = {"ResourceId": [resource_id]}
# Pass additional metadata as SSM parameters
for key, value in metadata.items():
if value:
parameters[key] = [str(value)]

response = ssm.start_automation_execution(
DocumentName=document_name,
Parameters=parameters,
)
execution_id = response["AutomationExecutionId"]
print(f"Started SSM Automation {document_name}, execution ID: {execution_id}")
return execution_id

def record_attempt(resource_id, remediation_type, execution_id):
"""
Write the remediation attempt to DynamoDB.
TTL is set to 24 hours so the table stays clean automatically.
"""
import time

history_table.put_item(
Item={
"resource_id": resource_id,
"timestamp": datetime.now(timezone.utc).isoformat(),
"remediation_type": remediation_type,
"ssm_execution_id": execution_id,
"ttl": int(time.time()) + (24 * 60 * 60), # 24-hour TTL
}
)

def notify_human(remediation_type, resource_id, attempts, metadata):
"""
When auto-remediation has failed too many times, escalate to a human
via SNS (which can route to PagerDuty, Slack, OpsGenie, etc.)
"""
message = {
"alert": "Auto-remediation escalation required",
"resource_id": resource_id,
"remediation_type": remediation_type,
"auto_attempts_made": attempts,
"metadata": metadata,
"action_required": "Manual investigation needed. Auto-remediation has been suppressed.",
}
sns.publish(
TopicArn=alert_topic_arn,
Subject=f"[AWS Auto-Heal] Escalation: {remediation_type} on {resource_id}",
Message=json.dumps(message, indent=2),
)