afterhours-shift-manager/scripts/cutover/recreate_holiday_schedules.py

#!/usr/bin/env python3
"""Recreate future holiday-* EventBridge Scheduler schedules in prod.

Reads outstanding holiday-activate-* / holiday-deactivate-* from mgmt and
creates the same names in prod targeting the prod router ARN and scheduler
role. Dry-run unless --execute.
"""

from __future__ import annotations

import argparse
import re
import sys
from datetime import datetime, timezone
from zoneinfo import ZoneInfo

import boto3
from botocore.exceptions import ClientError

SRC_ACCOUNT = "328440206208"
DST_ACCOUNT = "011934824531"
PROD_ROUTER_ARN = (
    "arn:aws:lambda:us-east-1:011934824531:function:afterhours-holiday-router"
)
PROD_ROLE_ARN = "arn:aws:iam::011934824531:role/tf-managed/afterhours-shift-manager-holiday-scheduler"
PREFIXES = ("holiday-activate-", "holiday-deactivate-")
_AT = re.compile(r"^at\((\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2})\)$")


def _client(profile: str, region: str):
    return boto3.Session(profile_name=profile, region_name=region).client("scheduler")


def _account(profile: str) -> str:
    return (
        boto3.Session(profile_name=profile)
        .client("sts")
        .get_caller_identity()["Account"]
    )


def schedule_when(detail: dict) -> datetime | None:
    """UTC instant the one-off schedule fires, or None if it cannot be parsed."""
    expr = (detail.get("ScheduleExpression") or "").strip()
    tzname = detail.get("ScheduleExpressionTimezone") or "America/New_York"
    match = _AT.match(expr)
    if match:
        naive = datetime.strptime(match.group(1), "%Y-%m-%dT%H:%M:%S")
        return naive.replace(tzinfo=ZoneInfo(tzname)).astimezone(timezone.utc)
    at = detail.get("EndDate") or detail.get("StartDate")
    if at is None:
        return None
    if at.tzinfo is None:
        return at.replace(tzinfo=timezone.utc)
    return at.astimezone(timezone.utc)


def _list_holiday(client):
    names = []
    token = None
    while True:
        kwargs = {"GroupName": "default"}
        if token:
            kwargs["NextToken"] = token
        resp = client.list_schedules(**kwargs)
        for item in resp.get("Schedules", []):
            name = item.get("Name", "")
            if name.startswith(PREFIXES):
                names.append(name)
        token = resp.get("NextToken")
        if not token:
            return names


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--src-profile", required=True)
    parser.add_argument("--dst-profile", required=True)
    parser.add_argument("--region", default="us-east-1")
    parser.add_argument("--execute", action="store_true")
    args = parser.parse_args()

    if _account(args.src_profile) != SRC_ACCOUNT:
        print("src profile is not mgmt", file=sys.stderr)
        return 2
    if _account(args.dst_profile) != DST_ACCOUNT:
        print("dst profile is not prod", file=sys.stderr)
        return 2

    src = _client(args.src_profile, args.region)
    dst = _client(args.dst_profile, args.region)
    now = datetime.now(timezone.utc)
    created = 0
    skipped = 0
    failed = 0

    for name in _list_holiday(src):
        detail = src.get_schedule(Name=name, GroupName="default")
        expr = detail.get("ScheduleExpression", "")
        tzname = detail.get("ScheduleExpressionTimezone", "America/New_York")
        when = schedule_when(detail)
        if when is not None and when < now:
            print(f"skip past {name} expr={expr}")
            skipped += 1
            continue
        payload = {
            "Name": name,
            "GroupName": "default",
            "ScheduleExpression": expr,
            "ScheduleExpressionTimezone": tzname,
            "FlexibleTimeWindow": {"Mode": "OFF"},
            "Target": {
                "Arn": PROD_ROUTER_ARN,
                "RoleArn": PROD_ROLE_ARN,
                "Input": detail.get("Target", {}).get("Input", ""),
            },
            "ActionAfterCompletion": detail.get("ActionAfterCompletion", "DELETE"),
        }
        if detail.get("EndDate"):
            payload["EndDate"] = detail["EndDate"]
        print(f"would create {name} expr={expr} tz={tzname}")
        if not args.execute:
            continue
        try:
            dst.create_schedule(**payload)
            created += 1
        except ClientError as exc:
            code = exc.response["Error"]["Code"]
            if code == "ConflictException":
                print(f"exists {name}")
            elif code == "ValidationException":
                print(
                    f"skip invalid {name}: {exc.response['Error'].get('Message', code)}"
                )
                skipped += 1
            else:
                print(f"failed {name}: {code}", file=sys.stderr)
                failed += 1

    print(
        f"created={created} skipped_past={skipped} failed={failed} execute={args.execute}"
    )
    if not args.execute:
        print("dry-run; pass --execute to CreateSchedule")
    return 1 if failed else 0


if __name__ == "__main__":
    raise SystemExit(main())