lambda_handler.py

from __future__ import annotations

import logging
import os
from collections.abc import Iterable, Mapping
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from decimal import Decimal, InvalidOperation
from typing import Any

try:
    import boto3
except ImportError:  # pragma: no cover - available in the deployed image
    boto3 = None

try:
    from botocore.exceptions import ClientError
except ImportError:  # pragma: no cover - available with boto3 in the deployed image
    class ClientError(Exception):
        response: dict[str, Any]

        def __init__(self, error_response: dict[str, Any], operation_name: str) -> None:
            super().__init__(operation_name)
            self.response = error_response


try:
    from lambdacron.lambda_task import CronLambdaTask
except ImportError:  # pragma: no cover - keeps local unit tests lightweight
    class CronLambdaTask:
        def __init__(self) -> None:
            self.logger = LOGGER

        def lambda_handler(self, event: Any, context: Any) -> dict[str, Any]:
            return self._perform_task(event, context)


logging.basicConfig(level=logging.INFO)
LOGGER = logging.getLogger(__name__)

DEFAULT_MAX_UPTIME_TAG_KEY = "max_hours_uptime"
DEFAULT_ENFORCEMENT_ACTION = "terminate"

OVERDUE_INSTANCES_RESULT_TYPE = "OVERDUE_INSTANCES"
STOP_REQUESTED_RESULT_TYPE = "STOP_REQUESTED"
TERMINATION_REQUESTED_RESULT_TYPE = "TERMINATION_REQUESTED"
INSTANCE_ACTION_ERRORS_RESULT_TYPE = "INSTANCE_ACTION_ERRORS"
INVALID_UPTIME_TAGS_RESULT_TYPE = "INVALID_UPTIME_TAGS"

RESULT_TYPE_BY_ACTION = {
    "notify": OVERDUE_INSTANCES_RESULT_TYPE,
    "stop": STOP_REQUESTED_RESULT_TYPE,
    "terminate": TERMINATION_REQUESTED_RESULT_TYPE,
}

NOTIFIABLE_STATES = frozenset({"pending", "running", "stopping", "stopped"})
STOPPABLE_STATES = frozenset({"running"})
TERMINATABLE_STATES = NOTIFIABLE_STATES


@dataclass(frozen=True)
class EnforcementConfig:
    tag_key: str
    action: str


def load_config(env: Mapping[str, str] | None = None) -> EnforcementConfig:
    source = env if env is not None else os.environ
    tag_key = source.get("MAX_UPTIME_TAG_KEY", DEFAULT_MAX_UPTIME_TAG_KEY).strip()
    action = source.get("ENFORCEMENT_ACTION", DEFAULT_ENFORCEMENT_ACTION).strip().lower()

    if not tag_key:
        raise ValueError("MAX_UPTIME_TAG_KEY must be non-empty.")

    if action not in RESULT_TYPE_BY_ACTION:
        supported_actions = ", ".join(sorted(RESULT_TYPE_BY_ACTION))
        raise ValueError(f"ENFORCEMENT_ACTION must be one of: {supported_actions}.")

    return EnforcementConfig(tag_key=tag_key, action=action)


def parse_max_hours(tag_value: str) -> Decimal:
    try:
        max_hours = Decimal(tag_value.strip())
    except (AttributeError, InvalidOperation) as exc:
        raise ValueError("Tag value must be a positive number of hours.") from exc

    if max_hours <= 0:
        raise ValueError("Tag value must be a positive number of hours.")

    return max_hours


def iter_tagged_instances(ec2_client: Any, tag_key: str) -> Iterable[dict[str, Any]]:
    filters = [{"Name": "tag-key", "Values": [tag_key]}]

    if hasattr(ec2_client, "get_paginator"):
        paginator = ec2_client.get_paginator("describe_instances")
        pages = paginator.paginate(Filters=filters)
    else:
        pages = [ec2_client.describe_instances(Filters=filters)]

    for page in pages:
        for reservation in page.get("Reservations", []):
            for instance in reservation.get("Instances", []):
                yield instance


def get_tag_value(tags: Iterable[Mapping[str, Any]] | None, key: str) -> str | None:
    for tag in tags or []:
        if tag.get("Key") == key:
            return tag.get("Value")
    return None


def get_instance_name(instance: Mapping[str, Any]) -> str | None:
    return get_tag_value(instance.get("Tags"), "Name")


def normalize_launch_time(value: datetime) -> datetime:
    if value.tzinfo is None:
        return value.replace(tzinfo=timezone.utc)
    return value.astimezone(timezone.utc)


def format_decimal(value: Decimal) -> str:
    normalized = format(value.normalize(), "f")
    return normalized.rstrip("0").rstrip(".") or "0"


def build_instance_payload(
    instance: Mapping[str, Any],
    *,
    tag_key: str,
    tag_value: str,
    max_hours: Decimal,
    now: datetime,
    action_taken: str,
    action_status: str,
    extra: Mapping[str, Any] | None = None,
) -> dict[str, Any]:
    launch_time = normalize_launch_time(instance["LaunchTime"])
    elapsed_hours = round((now - launch_time).total_seconds() / 3600, 2)
    deadline = launch_time + timedelta(hours=float(max_hours))
    state = (instance.get("State") or {}).get("Name", "unknown")

    payload = {
        "instance_id": instance["InstanceId"],
        "instance_name": get_instance_name(instance),
        "instance_state": state,
        "launch_time": launch_time.isoformat(),
        "deadline_time": deadline.isoformat(),
        "elapsed_hours": elapsed_hours,
        "max_hours_uptime": format_decimal(max_hours),
        "tag_key": tag_key,
        "tag_value": tag_value,
        "action_taken": action_taken,
        "action_status": action_status,
    }

    if extra:
        payload.update(extra)

    return payload


def build_result_payload(
    *,
    summary: str,
    config: EnforcementConfig,
    instances: list[dict[str, Any]],
    generated_at: datetime,
    requested_action: str,
) -> dict[str, Any]:
    return {
        "summary": summary,
        "configured_action": config.action,
        "requested_action": requested_action,
        "tag_key": config.tag_key,
        "instance_count": len(instances),
        "generated_at": generated_at.isoformat(),
        "instances": instances,
    }


def pluralize(noun: str, count: int) -> str:
    return noun if count == 1 else f"{noun}s"


def build_success_summary(action: str, count: int, tag_key: str) -> str:
    noun = pluralize("instance", count)

    if action == "notify":
        return f"Reported {count} overdue EC2 {noun} tagged with {tag_key}."
    if action == "stop":
        return f"Requested stop for {count} overdue EC2 {noun} tagged with {tag_key}."
    return f"Requested termination for {count} overdue EC2 {noun} tagged with {tag_key}."


def build_error_summary(action: str, count: int, tag_key: str) -> str:
    noun = pluralize("instance", count)
    return f"Failed to {action} {count} overdue EC2 {noun} tagged with {tag_key}."


def build_invalid_tag_summary(count: int, tag_key: str) -> str:
    noun = pluralize("instance", count)
    return f"Skipped {count} tagged EC2 {noun} because {tag_key} was not a positive number of hours."


def is_actionable_state(state: str, action: str) -> bool:
    if action == "notify":
        return state in NOTIFIABLE_STATES
    if action == "stop":
        return state in STOPPABLE_STATES
    return state in TERMINATABLE_STATES


def extract_client_error(exc: ClientError) -> dict[str, str]:
    error = getattr(exc, "response", {}).get("Error", {})
    return {
        "error_code": error.get("Code", exc.__class__.__name__),
        "error_message": error.get("Message", str(exc)),
    }


def run_enforcement(
    ec2_client: Any,
    *,
    config: EnforcementConfig,
    now: datetime | None = None,
) -> dict[str, dict[str, Any]]:
    evaluation_time = normalize_launch_time(now or datetime.now(timezone.utc))
    successful_instances: list[dict[str, Any]] = []
    invalid_tag_instances: list[dict[str, Any]] = []
    action_error_instances: list[dict[str, Any]] = []

    for instance in iter_tagged_instances(ec2_client, config.tag_key):
        tag_value = get_tag_value(instance.get("Tags"), config.tag_key)
        if tag_value is None:
            continue

        launch_time = normalize_launch_time(instance["LaunchTime"])
        elapsed_hours = Decimal((evaluation_time - launch_time).total_seconds()) / Decimal(3600)
        state = (instance.get("State") or {}).get("Name", "unknown")
        LOGGER.info(
            "Found candidate instance %s (%s) in state %s with uptime %.2f hours and %s=%s.",
            instance["InstanceId"],
            get_instance_name(instance) or "unnamed",
            state,
            float(elapsed_hours),
            config.tag_key,
            tag_value,
        )

        try:
            max_hours = parse_max_hours(tag_value)
        except ValueError as exc:
            invalid_tag_instances.append(
                {
                    "instance_id": instance["InstanceId"],
                    "instance_name": get_instance_name(instance),
                    "instance_state": (instance.get("State") or {}).get("Name", "unknown"),
                    "tag_key": config.tag_key,
                    "tag_value": tag_value,
                    "skip_reason": str(exc),
                }
            )
            continue

        if elapsed_hours < max_hours:
            continue

        if not is_actionable_state(state, config.action):
            LOGGER.info(
                "Skipping overdue instance %s in state %s for action %s.",
                instance["InstanceId"],
                state,
                config.action,
            )
            continue

        if config.action == "notify":
            successful_instances.append(
                build_instance_payload(
                    instance,
                    tag_key=config.tag_key,
                    tag_value=tag_value,
                    max_hours=max_hours,
                    now=evaluation_time,
                    action_taken="notify",
                    action_status="reported",
                )
            )
            continue

        action_fn = ec2_client.stop_instances if config.action == "stop" else ec2_client.terminate_instances

        try:
            action_fn(InstanceIds=[instance["InstanceId"]])
        except ClientError as exc:
            action_error_instances.append(
                build_instance_payload(
                    instance,
                    tag_key=config.tag_key,
                    tag_value=tag_value,
                    max_hours=max_hours,
                    now=evaluation_time,
                    action_taken=config.action,
                    action_status="failed",
                    extra=extract_client_error(exc),
                )
            )
            continue

        successful_instances.append(
            build_instance_payload(
                instance,
                tag_key=config.tag_key,
                tag_value=tag_value,
                max_hours=max_hours,
                now=evaluation_time,
                action_taken=config.action,
                action_status="requested",
            )
        )

    results: dict[str, dict[str, Any]] = {}

    if successful_instances:
        success_result_type = RESULT_TYPE_BY_ACTION[config.action]
        results[success_result_type] = build_result_payload(
            summary=build_success_summary(config.action, len(successful_instances), config.tag_key),
            config=config,
            instances=successful_instances,
            generated_at=evaluation_time,
            requested_action=config.action,
        )

    if invalid_tag_instances:
        results[INVALID_UPTIME_TAGS_RESULT_TYPE] = build_result_payload(
            summary=build_invalid_tag_summary(len(invalid_tag_instances), config.tag_key),
            config=config,
            instances=invalid_tag_instances,
            generated_at=evaluation_time,
            requested_action="skip",
        )

    if action_error_instances:
        results[INSTANCE_ACTION_ERRORS_RESULT_TYPE] = build_result_payload(
            summary=build_error_summary(config.action, len(action_error_instances), config.tag_key),
            config=config,
            instances=action_error_instances,
            generated_at=evaluation_time,
            requested_action=config.action,
        )

    return results


class EC2ZombieKillerTask(CronLambdaTask):
    def __init__(self, ec2_client: Any | None = None) -> None:
        super().__init__()
        self._ec2_client = ec2_client

    def _get_ec2_client(self) -> Any:
        if self._ec2_client is not None:
            return self._ec2_client

        if boto3 is None:  # pragma: no cover - only relevant outside the runtime image
            raise RuntimeError("boto3 is required to execute the EC2 zombie killer Lambda.")

        self._ec2_client = boto3.client("ec2")
        return self._ec2_client

    def _perform_task(self, event: Any, context: Any) -> dict[str, Any]:
        del event, context
        return run_enforcement(self._get_ec2_client(), config=load_config())


handler = EC2ZombieKillerTask().lambda_handler