# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 import argparse import json from argparse import ArgumentParser from typing import TYPE_CHECKING, Any, Literal, Optional, Union import boto3 import jmespath from instance_scheduler_cli import __version__ EVENT_SOURCE = "scheduler.cli" HELP_CMD_CREATE_PERIOD = "Creates a period" HELP_CMD_CREATE_SCHEDULE = "Creates a schedule" HELP_CMD_DELETE_SCHEDULE = "Deletes a schedule" HELP_CMD_DELETE_PERIOD = "Deletes a period" HELP_CMD_DESCRIBE_PERIODS = "Describes configured periods" HELP_CMD_DESCRIBE_SCHEDULES = "Describes configured schedules" HELP_CMD_SCHEDULE_DESCRIBE_USAGE = ( "Calculates periods and billing hours in which instances are running" ) HELP_CMD_UPDATE_PERIOD = "Updates a period" HELP_CMD_UPDATE_SCHEDULE = "Updates a schedule" HELP_ENDDATE = "End time of the period in format yyyymmdd, default is today" HELP_NAME_SCHEDULE = "Name of the schedule" HELP_PERIOD_BEGINTIME = "Begin time of the period in format hh:mm" HELP_PERIOD_DESCRIPTION = "Description for the period" HELP_PERIOD_ENDTIME = "End time of the period in format hh:mm" HELP_PERIOD_MONTH_DAYS = "Calendar monthdays of the period" HELP_PERIOD_MONTHS = "Months of the period" HELP_PERIOD_NAME = "Name of the period" HELP_PERIOD_WEEKDAYS = "Weekdays of the period" HELP_QUERY = "JMESPath query to transform or filter the result" HELP_REGION = "Region in which the Instance Scheduler stack is deployed" HELP_SCHEDULE_CLOUDWATCH_METRICS = "Enable CloudWatch metrics for this schedule" HELP_SCHEDULE_DESCRIPTION = "Description for the schedule." HELP_SCHEDULE_ENFORCED = "Enforce schedule state for instance." HELP_SCHEDULE_HIBERNATE = "Hibernate EC2 instances if possible when stopped." HELP_SCHEDULE_RETAIN_RUNNING = "Keep instances running at end of period if they were already running at start of period" HELP_SCHEDULE_NAME = "Name of the schedule" HELP_SCHEDULE_SSM_MAINTENANCE_WINDOW = ( "Name of SSM window in which EC2 instances are started" ) HELP_PARAM_TIMEZONE = "Timezone for schedule" HELP_SCHEDULE_OVERRIDE_STATUS = "Override status to keep instances in specified state." HELP_SCHEDULE_PERIODS = ( "List of the names of the periods in the schedule. Each period can specify an instance type by " "appending @<type> to the name of the period." ) HELP_SCHEDULE_KEEP_NEW = ( "Do not stop new instances if outside of a running period until end of next period" ) HELP_SCHEDULE_USE_MAIN = ( "Use prefered maintenace windows of RDS instances as a running period." ) HELP_STACK = "Name of the Instance Scheduler stack" HELP_PROFILE_NAME = ( " The name of a profile to use. If not given, then the default profile is used." ) HELP_STARTDATE = "Start time of the period in format yyyymmdd, default is today" HELP_SUB_COMMANDS = "Commands help" HELP_VALID_COMMANDS = "Valid subcommands" PROG_NAME = "scheduler-cli" VALUES_OVERRIDE_STATUS = ["stopped", "running"] PARAM_BEGINTIME = "--begintime" PARAM_DESCRIPTION = "--description" PARAM_ENDDATE = "--enddate" PARAM_ENDTIME = "--endtime" PARAM_ENFORCED = "--enforced" PARAM_HIBERNATE = "--hibernate" PARAM_RETAINED_RUNNING = "--retain-running" PARAM_METRICS = "--use-metrics" PARAM_MONTHDAYS = "--monthdays" PARAM_MONTHS = "--months" PARAM_OVERRIDE = "--override-status" PARAM_PERIODS = "--periods" PARAM_STARTDATE = "--startdate" PARAM_KEEP_NEW = "--do-not-stop-new-instances" PARAM_USE_MAIN = "--use-maintenance-window" PARAM_WEEKDAYS = "--weekdays" PARAM_TIMEZONE = "--timezone" PARAM_SSM_MAINTENCE_WINDOW = "--ssm-maintenance-window" PARAM_STACK = "--stack" PARAM_REGION = "--region" PARAM_QUERY = "--query" PARAM_PROFILE_NAME = "--profile-name" COMMON_PARAMS = [ s[2:].replace("-", "_") for s in [PARAM_QUERY, PARAM_REGION, PARAM_STACK, PARAM_PROFILE_NAME] ] + ["command"] CMD_CREATE_PERIOD = "create-period" CMD_CREATE_SCHEDULE = "create-schedule" CMD_DELETE_PERIOD = "delete-period" CMD_DELETE_SCHEDULE = "delete-schedule" CMD_DESCRIBE_PERIODS = "describe-periods" CMD_DESCRIBE_SCHEDULE_USAGE = "describe-schedule-usage" CMD_DESCRIBE_SCHEDULES = "describe-schedules" CMD_UPDATE_PERIOD = "update-period" CMD_UPDATE_SCHEDULE = "update-schedule" CMD_VERSION = "--version" PARAM_NAME = "--name" if TYPE_CHECKING: from mypy_boto3_cloudformation.client import CloudFormationClient from mypy_boto3_lambda.client import LambdaClient else: CloudFormationClient = object LambdaClient = object def _service_client( service: Union[Literal["cloudformation"], Literal["lambda"]], region: Optional[str] = None, profile_name: Optional[str] = None, ) -> Any: session = ( boto3.Session() if profile_name is None else boto3.Session(profile_name=profile_name) ) return session.client(service_name=service, region_name=region) def handle_command(args: Any, command: str) -> int: try: cloudformation_client: CloudFormationClient = _service_client( "cloudformation", region=args.region, profile_name=args.profile_name ) lambda_resource = cloudformation_client.describe_stack_resource( StackName=args.stack, LogicalResourceId="Main" )["StackResourceDetail"] lambda_client: LambdaClient = _service_client( "lambda", region=args.region, profile_name=args.profile_name ) event = { "source": EVENT_SOURCE, "action": command, "parameters": { a: getattr(args, a) for a in args.__dict__ if ( a not in COMMON_PARAMS and getattr(args, a) is not None and not hasattr(getattr(args, a), "__call__") ) }, } payload = str.encode(json.dumps(event)) lambda_name = lambda_resource["PhysicalResourceId"] # start lambda function resp = lambda_client.invoke( FunctionName=lambda_name, InvocationType="RequestResponse", LogType="None", Payload=payload, ) # read lambda response and load json lambda_response = resp["Payload"].read().decode("utf-8") result = json.loads(lambda_response) # Error if api raised an exception if "Error" in result: print(result["Error"].capitalize()) return 1 # perform transformation of output if args.query: result = jmespath.search(args.query, result) # print output as formatted json print(json.dumps(result, indent=3)) return 0 except Exception as ex: print(ex) return 1 def build_parser() -> ArgumentParser: def add_common_arguments(parser: ArgumentParser) -> None: parser.add_argument(PARAM_QUERY, PARAM_QUERY[1:3], help=HELP_QUERY) parser.add_argument(PARAM_REGION, PARAM_REGION[1:3], help=HELP_REGION) parser.add_argument( PARAM_STACK, PARAM_STACK[1:3], required=True, help=HELP_STACK ) parser.add_argument( PARAM_PROFILE_NAME, PARAM_PROFILE_NAME[1:3], required=False, help=HELP_PROFILE_NAME, ) def add_period_arguments(period_parser: ArgumentParser) -> None: period_parser.add_argument(PARAM_BEGINTIME, help=HELP_PERIOD_BEGINTIME) period_parser.add_argument(PARAM_DESCRIPTION, help=HELP_PERIOD_DESCRIPTION) period_parser.add_argument(PARAM_ENDTIME, help=HELP_PERIOD_ENDTIME) period_parser.add_argument(PARAM_MONTHDAYS, help=HELP_PERIOD_MONTH_DAYS) period_parser.add_argument(PARAM_MONTHS, help=HELP_PERIOD_MONTHS) period_parser.add_argument(PARAM_NAME, required=True, help=HELP_PERIOD_NAME) period_parser.add_argument(PARAM_WEEKDAYS, help=HELP_PERIOD_WEEKDAYS) def add_schedule_arguments(schedule_parser: ArgumentParser) -> None: schedule_parser.add_argument(PARAM_DESCRIPTION, help=HELP_SCHEDULE_DESCRIPTION) schedule_parser.add_argument(PARAM_TIMEZONE, help=HELP_PARAM_TIMEZONE) schedule_parser.add_argument(PARAM_NAME, required=True, help=HELP_SCHEDULE_NAME) schedule_parser.add_argument( PARAM_OVERRIDE, choices=VALUES_OVERRIDE_STATUS, help=HELP_SCHEDULE_OVERRIDE_STATUS, ) schedule_parser.add_argument( PARAM_PERIODS, type=str, help=HELP_SCHEDULE_PERIODS ) schedule_parser.add_argument( PARAM_KEEP_NEW, dest="stop_new_instances", action="store_false", help=HELP_SCHEDULE_KEEP_NEW, ) schedule_parser.add_argument( PARAM_USE_MAIN, default=False, dest="use_maintenance_window", action="store_true", help=HELP_SCHEDULE_USE_MAIN, ) schedule_parser.add_argument( PARAM_SSM_MAINTENCE_WINDOW, help=HELP_SCHEDULE_SSM_MAINTENANCE_WINDOW, type=str, ) schedule_parser.add_argument( PARAM_RETAINED_RUNNING, default=False, dest="retain_running", action="store_true", help=HELP_SCHEDULE_RETAIN_RUNNING, ) schedule_parser.add_argument( PARAM_ENFORCED, default=False, dest="enforced", action="store_true", help=HELP_SCHEDULE_ENFORCED, ) schedule_parser.add_argument( PARAM_HIBERNATE, default=False, dest="hibernate", action="store_true", help=HELP_SCHEDULE_HIBERNATE, ) schedule_parser.add_argument( PARAM_METRICS, default=False, dest="use-metrics", action="store_true", help=HELP_SCHEDULE_CLOUDWATCH_METRICS, ) def build_describe_schedules_parser() -> None: sub_parser = subparsers.add_parser( CMD_DESCRIBE_SCHEDULES, help=HELP_CMD_DESCRIBE_SCHEDULES ) sub_parser.add_argument(PARAM_NAME, help=HELP_NAME_SCHEDULE) add_common_arguments(sub_parser) sub_parser.set_defaults(func=handle_command, command=CMD_DESCRIBE_SCHEDULES) def build_describe_periods_parser() -> None: sub_parser = subparsers.add_parser( CMD_DESCRIBE_PERIODS, help=HELP_CMD_DESCRIBE_PERIODS ) sub_parser.add_argument(PARAM_NAME, help=HELP_PERIOD_NAME) add_common_arguments(sub_parser) sub_parser.set_defaults(func=handle_command, command=CMD_DESCRIBE_PERIODS) def build_create_period_parser() -> None: sub_parser = subparsers.add_parser( CMD_CREATE_PERIOD, help=HELP_CMD_CREATE_PERIOD ) add_period_arguments(sub_parser) add_common_arguments(sub_parser) sub_parser.set_defaults(func=handle_command, command=CMD_CREATE_PERIOD) def build_create_schedule_parser() -> None: sub_parser = subparsers.add_parser( CMD_CREATE_SCHEDULE, help=HELP_CMD_CREATE_SCHEDULE ) add_schedule_arguments(sub_parser) add_common_arguments(sub_parser) sub_parser.set_defaults(func=handle_command, command=CMD_CREATE_SCHEDULE) def build_update_period_parser() -> None: sub_parser = subparsers.add_parser( CMD_UPDATE_PERIOD, help=HELP_CMD_UPDATE_PERIOD ) add_period_arguments(sub_parser) add_common_arguments(sub_parser) sub_parser.set_defaults(func=handle_command, command=CMD_UPDATE_PERIOD) def build_update_schedule_parser() -> None: sub_parser = subparsers.add_parser( CMD_UPDATE_SCHEDULE, help=HELP_CMD_UPDATE_SCHEDULE ) add_schedule_arguments(sub_parser) add_common_arguments(sub_parser) sub_parser.set_defaults(func=handle_command, command=CMD_UPDATE_SCHEDULE) def build_delete_period_parser() -> None: sub_parser = subparsers.add_parser( CMD_DELETE_PERIOD, help=HELP_CMD_DELETE_PERIOD ) sub_parser.add_argument(PARAM_NAME, help=HELP_PERIOD_NAME) add_common_arguments(sub_parser) sub_parser.set_defaults(func=handle_command, command=CMD_DELETE_PERIOD) def build_delete_schedule_parser() -> None: sub_parser = subparsers.add_parser( CMD_DELETE_SCHEDULE, help=HELP_CMD_DELETE_SCHEDULE ) sub_parser.add_argument(PARAM_NAME, PARAM_NAME[1:3], help=HELP_SCHEDULE_NAME) add_common_arguments(sub_parser) sub_parser.set_defaults(func=handle_command, command=CMD_DELETE_SCHEDULE) def build_describe_schedule_usage_parser() -> None: sub_parser = subparsers.add_parser( CMD_DESCRIBE_SCHEDULE_USAGE, help=HELP_CMD_SCHEDULE_DESCRIBE_USAGE ) sub_parser.add_argument(PARAM_ENDDATE, help=HELP_ENDDATE) sub_parser.add_argument(PARAM_NAME, required=True, help=HELP_SCHEDULE_NAME) sub_parser.add_argument(PARAM_STARTDATE, help=HELP_STARTDATE) add_common_arguments(sub_parser) sub_parser.set_defaults( func=handle_command, command=CMD_DESCRIBE_SCHEDULE_USAGE ) new_parser = argparse.ArgumentParser(prog=PROG_NAME) new_parser.add_argument( CMD_VERSION, action="version", version=f"'%(prog)s {__version__}'" ) subparsers = new_parser.add_subparsers( help=HELP_SUB_COMMANDS, description=HELP_VALID_COMMANDS ) build_create_period_parser() build_create_schedule_parser() build_delete_period_parser() build_delete_schedule_parser() build_describe_periods_parser() build_describe_schedule_usage_parser() build_describe_schedules_parser() build_update_period_parser() build_update_schedule_parser() return new_parser