handler.py

import json
import os
import time
import uuid
from decimal import Decimal

import boto3

dynamodb = boto3.resource("dynamodb")
sfn = boto3.client("stepfunctions")
scheduler = boto3.client("scheduler")

SCENARIOS_TABLE = os.environ["SCENARIOS_TABLE"]
STATE_MACHINE_ARN = os.environ["STATE_MACHINE_ARN"]
SCHEDULER_ROLE_ARN = os.environ["SCHEDULER_ROLE_ARN"]
SCHEDULER_GROUP_NAME = os.environ["SCHEDULER_GROUP_NAME"]
REGIONAL_STATE_MACHINES = json.loads(os.environ["REGIONAL_STATE_MACHINES"])

table = dynamodb.Table(SCENARIOS_TABLE)


def _json_default(value):
    if isinstance(value, Decimal):
        return int(value) if value % 1 == 0 else float(value)
    raise TypeError(f"Object of type {type(value).__name__} is not JSON serializable")


def _response(status_code, body):
    return {
        "statusCode": status_code,
        "headers": {"Content-Type": "application/json"},
        "body": json.dumps(body, default=_json_default),
    }


def _now():
    return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())


def _create_scenario(event):
    body = json.loads(event.get("body") or "{}")
    name = body.get("name")
    script_key = body.get("scriptKey")
    if not name or not script_key:
        return _response(400, {"message": "name and scriptKey are required"})

    scenario_id = str(uuid.uuid4())
    item = {
        "scenarioId": scenario_id,
        "name": name,
        "description": body.get("description", ""),
        "scriptKey": script_key,
        "defaultTaskCount": int(body.get("defaultTaskCount", 1)),
        "regions": body.get("regions", []),
        "createdAt": _now(),
        "updatedAt": _now(),
    }
    table.put_item(Item=item)
    return _response(201, item)


def _list_scenarios(event):
    result = table.scan()
    return _response(200, {"scenarios": result.get("Items", [])})


def _get_scenario(scenario_id):
    result = table.get_item(Key={"scenarioId": scenario_id})
    item = result.get("Item")
    if not item:
        return _response(404, {"message": "scenario not found"})
    return _response(200, item)


def _update_scenario(event, scenario_id):
    existing = table.get_item(Key={"scenarioId": scenario_id}).get("Item")
    if not existing:
        return _response(404, {"message": "scenario not found"})

    body = json.loads(event.get("body") or "{}")
    updates = {
        "name": body.get("name", existing["name"]),
        "description": body.get("description", existing.get("description", "")),
        "scriptKey": body.get("scriptKey", existing["scriptKey"]),
        "defaultTaskCount": int(body.get("defaultTaskCount", existing.get("defaultTaskCount", 1))),
        "regions": body.get("regions", existing.get("regions", [])),
        "updatedAt": _now(),
    }
    table.update_item(
        Key={"scenarioId": scenario_id},
        UpdateExpression="SET #n = :name, description = :description, scriptKey = :scriptKey, "
        "defaultTaskCount = :defaultTaskCount, regions = :regions, updatedAt = :updatedAt",
        ExpressionAttributeNames={"#n": "name"},
        ExpressionAttributeValues={
            ":name": updates["name"],
            ":description": updates["description"],
            ":scriptKey": updates["scriptKey"],
            ":defaultTaskCount": updates["defaultTaskCount"],
            ":regions": updates["regions"],
            ":updatedAt": updates["updatedAt"],
        },
    )
    return _response(200, {**existing, **updates})


def _delete_scenario(scenario_id):
    table.delete_item(Key={"scenarioId": scenario_id})
    return _response(204, {})


def start_test(scenario_id, task_count=None, regions=None):
    """Start an execution for a scenario. Kept event-shape-independent so a
    future non-API trigger (e.g. EventBridge) can call it directly.

    task_count applies per selected region (matching the pre-multi-region
    single-region semantics unchanged) rather than being split across
    regions."""
    scenario = table.get_item(Key={"scenarioId": scenario_id}).get("Item")
    if not scenario:
        return None

    count = int(task_count or scenario.get("defaultTaskCount", 1))
    selected_regions = regions or scenario.get("regions") or [os.environ["AWS_REGION"]]

    region_specs = []
    for region in selected_regions:
        state_machine_arn = REGIONAL_STATE_MACHINES.get(region)
        if state_machine_arn is None:
            raise ValueError(f"unsupported region: {region}")
        region_specs.append(
            {
                "region": region,
                "stateMachineArn": state_machine_arn,
                "taskIndexes": list(range(count)),
            }
        )

    test_id = f"{scenario_id}-{int(time.time())}"

    sfn.start_execution(
        stateMachineArn=STATE_MACHINE_ARN,
        input=json.dumps(
            {
                "testId": test_id,
                "scenarioId": scenario_id,
                "scriptKey": scenario["scriptKey"],
                "taskCount": count,
                "regionSpecs": region_specs,
            }
        ),
    )
    return test_id


def _start_test_http(event, scenario_id):
    body = json.loads(event.get("body") or "{}")
    try:
        test_id = start_test(scenario_id, body.get("taskCount"), body.get("regions"))
    except ValueError as exc:
        return _response(400, {"message": str(exc)})
    if test_id is None:
        return _response(404, {"message": "scenario not found"})
    return _response(202, {"testId": test_id})


# --- schedules (EventBridge Scheduler) ---
# Schedule metadata isn't mirrored into DynamoDB; EventBridge Scheduler is
# the source of truth, filtered per-scenario via a "<scenarioId>--" name
# prefix (see _schedule_belongs_to_scenario).


def _schedule_belongs_to_scenario(scenario_id, schedule_id):
    return schedule_id.startswith(f"{scenario_id}--")


def _create_schedule(scenario_id, body, function_arn):
    scenario = table.get_item(Key={"scenarioId": scenario_id}).get("Item")
    if not scenario:
        return None

    expression = body.get("scheduleExpression")
    if not expression or not expression.startswith(("at(", "rate(", "cron(")):
        raise ValueError("scheduleExpression must be at(...), rate(...), or cron(...)")

    task_count = body.get("taskCount")
    target_input = {"scenarioId": scenario_id}
    if task_count is not None:
        target_input["taskCount"] = int(task_count)
    regions = body.get("regions")
    if regions is not None:
        target_input["regions"] = regions

    schedule_name = f"{scenario_id}--{uuid.uuid4()}"
    state = body.get("state", "ENABLED")
    result = scheduler.create_schedule(
        Name=schedule_name,
        GroupName=SCHEDULER_GROUP_NAME,
        ScheduleExpression=expression,
        ScheduleExpressionTimezone=body.get("timezone", "UTC"),
        State=state,
        FlexibleTimeWindow={"Mode": "OFF"},
        Target={
            "Arn": function_arn,
            "RoleArn": SCHEDULER_ROLE_ARN,
            "Input": json.dumps(target_input),
        },
    )
    return {
        "scheduleId": schedule_name,
        "scenarioId": scenario_id,
        "scheduleExpression": expression,
        "taskCount": task_count,
        "state": state,
        "arn": result["ScheduleArn"],
    }


def _create_schedule_http(event, scenario_id, context):
    body = json.loads(event.get("body") or "{}")
    try:
        result = _create_schedule(scenario_id, body, context.invoked_function_arn)
    except ValueError as exc:
        return _response(400, {"message": str(exc)})
    if result is None:
        return _response(404, {"message": "scenario not found"})
    return _response(201, result)


def _list_schedules_http(scenario_id):
    result = scheduler.list_schedules(GroupName=SCHEDULER_GROUP_NAME, NamePrefix=f"{scenario_id}--")
    schedules = [
        {
            "scheduleId": s["Name"],
            "state": s.get("State"),
            "creationDate": s["CreationDate"].isoformat() if s.get("CreationDate") else None,
            "lastModified": s["LastModificationDate"].isoformat() if s.get("LastModificationDate") else None,
        }
        for s in result.get("Schedules", [])
    ]
    return _response(200, {"schedules": schedules})


def _get_schedule_http(scenario_id, schedule_id):
    if not _schedule_belongs_to_scenario(scenario_id, schedule_id):
        return _response(404, {"message": "schedule not found"})
    try:
        resp = scheduler.get_schedule(Name=schedule_id, GroupName=SCHEDULER_GROUP_NAME)
    except scheduler.exceptions.ResourceNotFoundException:
        return _response(404, {"message": "schedule not found"})

    target_input = json.loads(resp.get("Target", {}).get("Input", "{}"))
    return _response(
        200,
        {
            "scheduleId": resp["Name"],
            "scenarioId": scenario_id,
            "scheduleExpression": resp.get("ScheduleExpression"),
            "taskCount": target_input.get("taskCount"),
            "state": resp.get("State"),
            "arn": resp.get("Arn"),
            "flexibleTimeWindow": resp.get("FlexibleTimeWindow"),
        },
    )


def _delete_schedule_http(scenario_id, schedule_id):
    if not _schedule_belongs_to_scenario(scenario_id, schedule_id):
        return _response(404, {"message": "schedule not found"})
    try:
        scheduler.delete_schedule(Name=schedule_id, GroupName=SCHEDULER_GROUP_NAME)
    except scheduler.exceptions.ResourceNotFoundException:
        return _response(404, {"message": "schedule not found"})
    return _response(204, {})


def _handle_scheduled_invocation(event, context):
    # EventBridge Scheduler's Lambda target invokes with a bare payload
    # (no httpMethod/resource envelope) — exceptions intentionally propagate
    # here so Scheduler's own retry policy engages, unlike the HTTP branch
    # below which must always return a response.
    test_id = start_test(event["scenarioId"], event.get("taskCount"), event.get("regions"))
    if test_id is None:
        raise ValueError(f"scenario {event.get('scenarioId')} not found")
    return {"testId": test_id}


def lambda_handler(event, context):
    if "httpMethod" not in event:
        return _handle_scheduled_invocation(event, context)

    method = event["httpMethod"]
    resource = event["resource"]
    path_params = event.get("pathParameters") or {}
    scenario_id = path_params.get("id")
    schedule_id = path_params.get("scheduleId")

    try:
        if resource == "/scenarios" and method == "POST":
            return _create_scenario(event)
        if resource == "/scenarios" and method == "GET":
            return _list_scenarios(event)
        if resource == "/scenarios/{id}" and method == "GET":
            return _get_scenario(scenario_id)
        if resource == "/scenarios/{id}" and method == "PUT":
            return _update_scenario(event, scenario_id)
        if resource == "/scenarios/{id}" and method == "DELETE":
            return _delete_scenario(scenario_id)
        if resource == "/scenarios/{id}/start" and method == "POST":
            return _start_test_http(event, scenario_id)
        if resource == "/scenarios/{id}/schedules" and method == "POST":
            return _create_schedule_http(event, scenario_id, context)
        if resource == "/scenarios/{id}/schedules" and method == "GET":
            return _list_schedules_http(scenario_id)
        if resource == "/scenarios/{id}/schedules/{scheduleId}" and method == "GET":
            return _get_schedule_http(scenario_id, schedule_id)
        if resource == "/scenarios/{id}/schedules/{scheduleId}" and method == "DELETE":
            return _delete_schedule_http(scenario_id, schedule_id)
        return _response(404, {"message": "not found"})
    except Exception as exc:  # noqa: BLE001 - surface as a 500 with a message
        return _response(500, {"message": str(exc)})