War Story: We Survived a Key Leak Using AWS KMS Automatic Rotation

package main

import (
    "context"
    "encoding/json"
    "fmt"
    "log"
    "net/http"
    "os"
    "strings"
    "time"

    "github.com/aws/aws-sdk-go/aws"
    "github.com/aws/aws-sdk-go/aws/session"
    "github.com/aws/aws-sdk-go/service/kms"
    "github.com/google/go-github/v50/github"
)

const (
    kmsKeyPattern = `arn:aws:kms:[a-z0-9-]+:\d+:key/[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}`
    leakAlertSlackChannel = "#security-incidents"
)

// leakedKeyEvent represents a validated KMS key leak found in a GitHub push
type leakedKeyEvent struct {
    KeyARN   string    `json:"key_arn"`
    Repo     string    `json:"repo"`
    CommitSHA string  `json:"commit_sha"`
    DetectedAt time.Time `json:"detected_at"`
}

// kmsRotator handles KMS key rotation operations
type kmsRotator struct {
    client *kms.KMS
}

// newKMSRotator initializes a KMS client with the default AWS session
func newKMSRotator(region string) (*kmsRotator, error) {
    sess, err := session.NewSession(&aws.Config{
        Region: aws.String(region),
    })
    if err != nil {
        return nil, fmt.Errorf("failed to create AWS session: %w", err)
    }
    return &kmsRotator{client: kms.New(sess)}, nil
}

// validateKey checks if a KMS key ARN exists and is active
func (r *kmsRotator) validateKey(ctx context.Context, keyARN string) (bool, error) {
    // Extract key ID from ARN for DescribeKey call
    parts := strings.Split(keyARN, "/")
    if len(parts) != 2 {
        return false, fmt.Errorf("invalid key ARN format: %s", keyARN)
    }
    keyID := parts[1]

    _, err := r.client.DescribeKeyWithContext(ctx, &kms.DescribeKeyInput{
        KeyId: aws.String(keyID),
    })
    if err != nil {
        // Key doesn't exist or is inaccessible
        return false, nil
    }
    return true, nil
}

// triggerRotation initiates automatic rotation for a validated KMS key
func (r *kmsRotator) triggerRotation(ctx context.Context, keyARN string) error {
    parts := strings.Split(keyARN, "/")
    if len(parts) != 2 {
        return fmt.Errorf("invalid key ARN format: %s", keyARN)
    }
    keyID := parts[1]

    _, err := r.client.EnableKeyRotationWithContext(ctx, &kms.EnableKeyRotationInput{
        KeyId: aws.String(keyID),
    })
    if err != nil {
        return fmt.Errorf("failed to enable rotation for key %s: %w", keyID, err)
    }

    log.Printf("Successfully triggered rotation for KMS key %s", keyID)
    return nil
}

// handleGitHubWebhook processes push events from GitHub, scans for leaked KMS keys
func handleGitHubWebhook(w http.ResponseWriter, r *http.Request) {
    ctx := r.Context()
    rotator, err := newKMSRotator(os.Getenv("AWS_REGION"))
    if err != nil {
        log.Printf("Failed to initialize KMS rotator: %v", err)
        http.Error(w, "Internal server error", http.StatusInternalServerError)
        return
    }

    // Verify webhook signature (simplified for example; use proper HMAC verification in prod)
    payload, err := github.ValidatePayload(r, []byte(os.Getenv("GITHUB_WEBHOOK_SECRET")))
    if err != nil {
        log.Printf("Invalid webhook signature: %v", err)
        http.Error(w, "Unauthorized", http.StatusUnauthorized)
        return
    }
    defer r.Body.Close()

    var pushEvent github.PushEvent
    if err := json.Unmarshal(payload, &pushEvent); err != nil {
        log.Printf("Failed to unmarshal push event: %v", err)
        http.Error(w, "Bad request", http.StatusBadRequest)
        return
    }

    // Scan commit diffs for KMS key ARNs
    // In production, this would scan full file contents, not just diffs
    for _, commit := range pushEvent.Commits {
        for _, added := range commit.Added {
            // Check if added file contains KMS key ARN pattern
            // Simplified: check filename for demo; real impl uses regex scan of file content
            if strings.Contains(added, "arn:aws:kms") {
                // Extract key ARN using regex (simplified for example)
                keyARN := extractKeyARN(added)
                if keyARN == "" {
                    continue
                }

                // Validate key exists and is active
                valid, err := rotator.validateKey(ctx, keyARN)
                if err != nil {
                    log.Printf("Failed to validate key %s: %v", keyARN, err)
                    continue
                }
                if !valid {
                    log.Printf("Key %s is invalid or inaccessible, skipping", keyARN)
                    continue
                }

                // Trigger rotation immediately
                if err := rotator.triggerRotation(ctx, keyARN); err != nil {
                    log.Printf("Failed to rotate key %s: %v", keyARN, err)
                    // Send Slack alert for manual intervention
                    sendSlackAlert(ctx, leakedKeyEvent{
                        KeyARN:    keyARN,
                        Repo:      pushEvent.Repo.GetFullName(),
                        CommitSHA: commit.GetID(),
                        DetectedAt: time.Now(),
                    })
                    continue
                }

                log.Printf("Mitigated leak for key %s in repo %s", keyARN, pushEvent.Repo.GetFullName())
            }
        }
    }

    w.WriteHeader(http.StatusOK)
}

// extractKeyARN is a placeholder for real regex extraction; uses strings for demo
func extractKeyARN(content string) string {
    // In production, use regexp.MustCompile(kmsKeyPattern).FindString(content)
    if strings.Contains(content, "arn:aws:kms") {
        // Simplified extraction: return first matching substring
        start := strings.Index(content, "arn:aws:kms")
        if start == -1 {
            return ""
        }
        end := start + 100 // Simplified; real impl finds end of ARN
        return content[start:end]
    }
    return ""
}

// sendSlackAlert sends a critical alert to the security Slack channel
func sendSlackAlert(ctx context.Context, event leakedKeyEvent) {
    // In production, use slack-go/slack to send real messages
    log.Printf("CRITICAL: Leaked KMS key %s detected in repo %s, commit %s", event.KeyARN, event.Repo, event.CommitSHA)
}

func main() {
    http.HandleFunc("/github-webhook", handleGitHubWebhook)
    port := os.Getenv("PORT")
    if port == "" {
        port = "8080"
    }
    log.Printf("Starting webhook server on port %s", port)
    if err := http.ListenAndServe(":"+port, nil); err != nil {
        log.Fatalf("Failed to start server: %v", err)
    }
}

terraform {
  required_version = ">= 1.3.0"
  required_providers {
    aws = {
      source  = "hashicorp/aws"
      version = "~> 5.0"
    }
  }
}

provider "aws" {
  region = var.aws_region
}

variable "aws_region" {
  type        = string
  description = "AWS region to deploy KMS resources"
  default     = "us-east-1"
}

variable "environment" {
  type        = string
  description = "Deployment environment (prod, staging, dev)"
  validation {
    condition     = contains(["prod", "staging", "dev"], var.environment)
    error_message = "Environment must be one of: prod, staging, dev."
  }
}

variable "kms_key_deletion_window" {
  type        = number
  description = "Waiting period before KMS key deletion (7-30 days)"
  default     = 30
  validation {
    condition     = var.kms_key_deletion_window >= 7 && var.kms_key_deletion_window <= 30
    error_message = "Deletion window must be between 7 and 30 days."
  }
}

# Customer-managed KMS key for production user data encryption
resource "aws_kms_key" "user_data" {
  description             = "KMS key for encrypting production user data - ${var.environment}"
  deletion_window_in_days = var.kms_key_deletion_window
  enable_key_rotation     = true # Enable automatic annual rotation
  key_usage               = "ENCRYPT_DECRYPT"
  customer_master_key_spec = "SYMMETRIC_DEFAULT"

  # Policy restricts key usage to production services only
  policy = jsonencode({
    Version = "2012-10-17"
    Statement = [
      {
        Sid    = "AllowRootAccountFullAccess"
        Effect = "Allow"
        Principal = {
          AWS = "arn:aws:iam::${data.aws_caller_identity.current.account_id}:root"
        }
        Action   = "kms:*"
        Resource = "*"
      },
      {
        Sid    = "AllowProductionServicesDecrypt"
        Effect = "Allow"
        Principal = {
          AWS = aws_iam_role.production_services.arn
        }
        Action = [
          "kms:Decrypt",
          "kms:DescribeKey",
          "kms:Encrypt",
          "kms:GenerateDataKey*",
          "kms:ReEncrypt*"
        ]
        Resource = "*"
        Condition = {
          StringEquals = {
            "kms:EncryptionContext:environment" = var.environment
          }
        }
      },
      {
        Sid    = "DenyKeyDeletionForNonRoot"
        Effect = "Deny"
        Principal = "*"
        Action = "kms:ScheduleKeyDeletion"
        Resource = "*"
        Condition = {
          StringNotEquals = {
            "aws:PrincipalArn" = "arn:aws:iam::${data.aws_caller_identity.current.account_id}:root"
          }
        }
      }
    ]
  })

  tags = {
    Environment = var.environment
    Purpose     = "User Data Encryption"
    ManagedBy   = "Terraform"
  }
}

# Alias for the KMS key to simplify reference in application code
resource "aws_kms_alias" "user_data" {
  name          = "alias/user-data-${var.environment}"
  target_key_id = aws_kms_key.user_data.key_id
}

# IAM role for production services that need to access the KMS key
resource "aws_iam_role" "production_services" {
  name = "prod-services-kms-role-${var.environment}"

  assume_role_policy = jsonencode({
    Version = "2012-10-17"
    Statement = [
      {
        Action = "sts:AssumeRole"
        Effect = "Allow"
        Principal = {
          Service = ["ec2.amazonaws.com", "lambda.amazonaws.com"]
        }
      }
    ]
  })

  tags = {
    Environment = var.environment
    Purpose     = "KMS Access for Production Services"
  }
}

# CloudWatch alarm for KMS key rotation failures
resource "aws_cloudwatch_metric_alarm" "kms_rotation_failure" {
  alarm_name          = "kms-key-rotation-failure-${var.environment}"
  comparison_operator = "GreaterThanThreshold"
  evaluation_periods  = 1
  metric_name         = "KeyRotationFailureCount"
  namespace           = "AWS/KMS"
  period              = 300
  statistic           = "Sum"
  threshold           = 0
  alarm_description   = "Triggers when KMS key rotation fails for ${var.environment} user data key"
  alarm_actions       = [aws_sns_topic.security_alerts.arn]

  dimensions = {
    KeyId = aws_kms_key.user_data.key_id
  }

  tags = {
    Environment = var.environment
    Purpose     = "KMS Rotation Monitoring"
  }
}

# SNS topic for security alerts
resource "aws_sns_topic" "security_alerts" {
  name = "security-alerts-${var.environment}"
}

data "aws_caller_identity" "current" {}

output "kms_key_arn" {
  description = "ARN of the provisioned KMS key"
  value       = aws_kms_key.user_data.arn
}

output "kms_key_alias" {
  description = "Alias of the provisioned KMS key"
  value       = aws_kms_alias.user_data.name
}

import boto3
import json
import csv
from datetime import datetime
from botocore.exceptions import ClientError, NoCredentialsError
from typing import Dict, List, Optional

# Configuration
REPORT_OUTPUT_PATH = "kms_rotation_audit_report.csv"
ROTATION_ENABLED_REQUIRED = True
SUPPORTED_REGIONS = [
    "us-east-1", "us-east-2", "us-west-1", "us-west-2",
    "eu-west-1", "eu-west-2", "eu-central-1", "ap-southeast-1", "ap-southeast-2"
]

def get_kms_client(region: str) -> boto3.client:
    """Initialize a KMS client for the specified region with error handling."""
    try:
        return boto3.client("kms", region_name=region)
    except NoCredentialsError:
        raise RuntimeError(f"No AWS credentials found for region {region}")
    except ClientError as e:
        raise RuntimeError(f"Failed to create KMS client for {region}: {e}")

def list_customer_managed_keys(kms_client: boto3.client) -> List[Dict]:
    """List all customer-managed KMS keys in the region, handling pagination."""
    keys = []
    paginator = kms_client.get_paginator("list_keys")
    try:
        for page in paginator.paginate():
            for key in page.get("Keys", []):
                # Skip AWS managed keys (they have a different ARN pattern)
                key_arn = key.get("KeyArn", "")
                if "aws:managed" in key_arn.lower():
                    continue
                keys.append(key)
    except ClientError as e:
        print(f"Failed to list keys: {e}")
        return []
    return keys

def get_key_rotation_status(kms_client: boto3.client, key_id: str) -> Optional[bool]:
    """Check if automatic rotation is enabled for a KMS key."""
    try:
        response = kms_client.get_key_rotation_status(KeyId=key_id)
        return response.get("KeyRotationEnabled", False)
    except ClientError as e:
        if e.response["Error"]["Code"] == "AccessDeniedException":
            print(f"Access denied for key {key_id}")
        else:
            print(f"Failed to get rotation status for key {key_id}: {e}")
        return None

def get_key_metadata(kms_client: boto3.client, key_id: str) -> Optional[Dict]:
    """Retrieve full metadata for a KMS key."""
    try:
        response = kms_client.describe_key(KeyId=key_id)
        return response.get("KeyMetadata", {})
    except ClientError as e:
        print(f"Failed to describe key {key_id}: {e}")
        return None

def audit_kms_rotation() -> None:
    """Main audit function: scans all regions, checks rotation status, generates report."""
    audit_results = []
    total_keys = 0
    compliant_keys = 0
    non_compliant_keys = 0
    error_keys = 0

    print(f"Starting KMS rotation audit across {len(SUPPORTED_REGIONS)} regions...")

    for region in SUPPORTED_REGIONS:
        print(f"Scanning region: {region}")
        try:
            kms_client = get_kms_client(region)
        except RuntimeError as e:
            print(f"Skipping region {region}: {e}")
            continue

        # List all customer-managed keys in the region
        region_keys = list_customer_managed_keys(kms_client)
        total_keys += len(region_keys)
        print(f"Found {len(region_keys)} customer-managed keys in {region}")

        for key in region_keys:
            key_id = key.get("KeyId")
            key_arn = key.get("KeyArn")

            # Get key metadata
            metadata = get_key_metadata(kms_client, key_id)
            if not metadata:
                error_keys += 1
                audit_results.append({
                    "key_id": key_id,
                    "key_arn": key_arn,
                    "region": region,
                    "rotation_enabled": "ERROR",
                    "key_state": "UNKNOWN",
                    "last_rotated": "UNKNOWN",
                    "compliant": "NO"
                })
                continue

            # Get rotation status
            rotation_enabled = get_key_rotation_status(kms_client, key_id)
            rotation_status_str = str(rotation_enabled) if rotation_enabled is not None else "ERROR"

            # Check compliance
            compliant = "YES" if (rotation_enabled is True and ROTATION_ENABLED_REQUIRED) else "NO"
            if compliant == "YES":
                compliant_keys += 1
            else:
                non_compliant_keys += 1

            # Get last rotation time if available
            last_rotated = metadata.get("LastRotationDate", "NEVER")
            if last_rotated != "NEVER":
                last_rotated = last_rotated.strftime("%Y-%m-%d %H:%M:%S")

            audit_results.append({
                "key_id": key_id,
                "key_arn": key_arn,
                "region": region,
                "rotation_enabled": rotation_status_str,
                "key_state": metadata.get("KeyState", "UNKNOWN"),
                "last_rotated": last_rotated,
                "compliant": compliant
            })

    # Generate CSV report
    if audit_results:
        try:
            with open(REPORT_OUTPUT_PATH, "w", newline="") as csvfile:
                fieldnames = [
                    "key_id", "key_arn", "region", "rotation_enabled",
                    "key_state", "last_rotated", "compliant"
                ]
                writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
                writer.writeheader()
                writer.writerows(audit_results)
            print(f"Audit report saved to {REPORT_OUTPUT_PATH}")
        except IOError as e:
            print(f"Failed to write report: {e}")
    else:
        print("No keys found to audit.")

    # Print summary
    print("\n=== Audit Summary ===")
    print(f"Total customer-managed keys scanned: {total_keys}")
    print(f"Compliant keys (rotation enabled): {compliant_keys}")
    print(f"Non-compliant keys (rotation disabled): {non_compliant_keys}")
    print(f"Keys with errors (access denied/invalid): {error_keys}")
    print(f"Compliance rate: {(compliant_keys / total_keys * 100) if total_keys > 0 else 0:.2f}%")

if __name__ == "__main__":
    try:
        audit_kms_rotation()
    except KeyboardInterrupt:
        print("\nAudit interrupted by user.")
    except Exception as e:
        print(f"Audit failed with unexpected error: {e}")

aws kms enable-key-rotation --key-id arn:aws:kms:us-east-1:123456789012:key/1234abcd-12ab-34cd-56ef-1234567890ab

checkov -d terraform/ --framework terraform --check CKV_AWS_42

import aws_encryption_sdk
from aws_encryption_sdk import CommitmentPolicy

client = aws_encryption_sdk.EncryptionSDKClient(
    commitment_policy=CommitmentPolicy.REQUIRE_ENCRYPT_REQUIRE_DECRYPT
)
kms_key_arn = "arn:aws:kms:us-east-1:123456789012:key/1234abcd-12ab-34cd-56ef-1234567890ab"
with open("plaintext.txt", "rb") as f:
    plaintext = f.read()
encrypted_data, _ = client.encrypt(
    source=plaintext,
    key_ids=[kms_key_arn]
)