Skip to content

Commit

Permalink
Refactored to work with new methodology that scales.
Browse files Browse the repository at this point in the history
  • Loading branch information
@mk-amz
mk-amz committed Mar 18, 2024
1 parent e55037e commit a67f6c5
Show file tree
Hide file tree
Showing 6 changed files with 208 additions and 97 deletions.
100 changes: 73 additions & 27 deletions aws_sra_examples/solutions/patch_mgmt/lambda/src/app.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,9 @@
LOGGER.setLevel(log_level)

# Initialize the helper. `sleep_on_delete` allows time for the CloudWatch Logs to get captured.
helper = CfnResource(json_logging=True, log_level=log_level, boto_level="CRITICAL", sleep_on_delete=120)
helper = CfnResource(
json_logging=True, log_level=log_level, boto_level="CRITICAL", sleep_on_delete=120
)

# Global variables
UNEXPECTED = "Unexpected!"
Expand All @@ -58,7 +60,9 @@ class MaintInfo(TypedDict):
window_ids: list


def create_maintenance_window(params: dict, account_id: str, regions: list) -> MaintInfo:
def create_maintenance_window(
params: dict, account_id: str, regions: list
) -> MaintInfo:
"""Create a maintenance window.
Args:
Expand All @@ -76,16 +80,22 @@ def create_maintenance_window(params: dict, account_id: str, regions: list) -> M
)
window_ids = []
for region in regions:
LOGGER.info(f"Setting up Default Host Management and Creating a Maint Window {account_id} {region}")
LOGGER.info(
f"Setting up Default Host Management and Creating a Maint Window {account_id} {region}"
)
ssmclient = session.client("ssm", region_name=region, config=boto3_config)
ssmclient.update_service_setting(
SettingId="/ssm/managed-instance/default-ec2-instance-management-role",
SettingValue="service-role/AWSSystemsManagerDefaultEC2InstanceManagementRole",
)
maintenance_window_name = params.get("MAINTENANCE_WINDOW_NAME", "")
maintenance_window_description = params.get("MAINTENANCE_WINDOW_DESCRIPTION", "")
maintenance_window_description = params.get(
"MAINTENANCE_WINDOW_DESCRIPTION", ""
)
maintenance_window_schedule = params.get("MAINTENANCE_WINDOW_SCHEDULE", "")
maintenance_window_duration = int(params.get("MAINTENANCE_WINDOW_DURATION", 120))
maintenance_window_duration = int(
params.get("MAINTENANCE_WINDOW_DURATION", 120)
)
maintenance_window_cutoff = int(params.get("MAINTENANCE_WINDOW_CUTOFF", 0))
maintenance_window_timezone = params.get("MAINTENANCE_WINDOW_TIMEZONE", "")

Expand All @@ -97,6 +107,7 @@ def create_maintenance_window(params: dict, account_id: str, regions: list) -> M
Cutoff=maintenance_window_cutoff,
ScheduleTimezone=maintenance_window_timezone,
AllowUnassociatedTargets=False,
Tags=[{"Key": "createdBy", "Value": "SRA_Patch_Management"}]
)
window_ids.append(
{
Expand All @@ -109,7 +120,9 @@ def create_maintenance_window(params: dict, account_id: str, regions: list) -> M
return {"window_ids": window_ids}


def define_maintenance_window_targets(params: dict, window_id_response: list, account_id: str) -> list[dict[str, Any]]:
def define_maintenance_window_targets(
params: dict, window_id_response: list, account_id: str
) -> list[dict[str, Any]]:
"""Define Maintenance Window Targets.
Args:
Expand All @@ -128,7 +141,9 @@ def define_maintenance_window_targets(params: dict, window_id_response: list, ac
window_targets = []
for response in window_id_response:
LOGGER.info(f"Maintenance Window Targets {response['region']}")
ssmclient = session.client("ssm", region_name=response["region"], config=boto3_config)
ssmclient = session.client(
"ssm", region_name=response["region"], config=boto3_config
)

# Target Args for SSM Update
target_name = params.get("TARGET_NAME", "")
Expand Down Expand Up @@ -162,7 +177,12 @@ def define_maintenance_window_targets(params: dict, window_id_response: list, ac
return window_targets


def define_maintenance_window_tasks(params: dict, window_id_response: list, window_target_response: list, account_id: str) -> list[dict[str, Any]]:
def define_maintenance_window_tasks(
params: dict,
window_id_response: list,
window_target_response: list,
account_id: str,
) -> list[dict[str, Any]]:
"""Define maintenance window targets.
Args:
Expand All @@ -183,7 +203,9 @@ def define_maintenance_window_tasks(params: dict, window_id_response: list, wind
for response in window_id_response:
LOGGER.info(f"Maintenance Window Tasks in {response['region']}")
LOGGER.info(response)
ssmclient = session.client("ssm", region_name=response["region"], config=boto3_config)
ssmclient = session.client(
"ssm", region_name=response["region"], config=boto3_config
)
# Task Args for SSM Update
task_name = params.get("TASK_NAME", "")
task_description = params.get("TASK_DESCRIPTION", "")
Expand All @@ -207,7 +229,7 @@ def define_maintenance_window_tasks(params: dict, window_id_response: list, wind
TaskArn=task_run_command,
TaskType="RUN_COMMAND",
Priority=1,
ServiceRoleArn=f"arn:aws:iam::{account_id}:role/AmazonSSMAutomationRole",
ServiceRoleArn=f"arn:aws:iam::{account_id}:role/sra-patch-mgmt-automation",
CutoffBehavior="CONTINUE_TASK",
MaxConcurrency="100",
MaxErrors="1",
Expand All @@ -233,7 +255,9 @@ def define_maintenance_window_tasks(params: dict, window_id_response: list, wind
return window_ids


def parameter_pattern_validator(parameter_name: str, parameter_value: str, pattern: str) -> None:
def parameter_pattern_validator(
parameter_name: str, parameter_value: str, pattern: str
) -> None:
"""Validate CloudFormation Custom Resource Parameters.
Args:
Expand All @@ -245,7 +269,9 @@ def parameter_pattern_validator(parameter_name: str, parameter_value: str, patte
ValueError: Parameter does not follow the allowed pattern
"""
if not re.match(pattern, parameter_value):
raise ValueError(f"'{parameter_name}' parameter with value of '{parameter_value}' does not follow the allowed pattern: {pattern}.")
raise ValueError(
f"'{parameter_name}' parameter with value of '{parameter_value}' does not follow the allowed pattern: {pattern}."
)


def get_validated_parameters(
Expand Down Expand Up @@ -318,7 +344,9 @@ def get_validated_parameters(
params.get("MAINTENANCE_WINDOW_TIMEZONE", ""),
pattern=r"^[\w\/+=,.@-]{1,64}$",
)
parameter_pattern_validator("TASK_NAME", params.get("TASK_NAME", ""), pattern=r"^[\w\s+=,.@-]{1,64}$")
parameter_pattern_validator(
"TASK_NAME", params.get("TASK_NAME", ""), pattern=r"^[\w\s+=,.@-]{1,64}$"
)
parameter_pattern_validator(
"TASK_DESCRIPTION",
params.get("TASK_DESCRIPTION", ""),
Expand All @@ -329,7 +357,9 @@ def get_validated_parameters(
params.get("TASK_RUN_COMMAND", ""),
pattern=r"^[\w\s+=,.@-]{1,64}$",
)
parameter_pattern_validator("TARGET_NAME", params.get("TARGET_NAME", ""), pattern=r"^[\w\s+=,.@-]{1,64}$")
parameter_pattern_validator(
"TARGET_NAME", params.get("TARGET_NAME", ""), pattern=r"^[\w\s+=,.@-]{1,64}$"
)
parameter_pattern_validator(
"TARGET_DESCRIPTION",
params.get("TARGET_DESCRIPTION", ""),
Expand Down Expand Up @@ -364,27 +394,42 @@ def process_create_update_event(params: dict, regions: list) -> Dict:
Returns:
Dict: Dictionary of Window IDs, Targets, and Tasks
"""
account_ids = common.get_account_ids([], params["DELEGATED_ADMIN_ACCOUNT_ID"]) # they updated the stack and want us to remove things.
if (params.get("DISABLE_PATCHMGMT", "false")).lower() in "true" and params["action"] == "Update":
account_ids = common.get_account_ids(
[], params["DELEGATED_ADMIN_ACCOUNT_ID"]
) # they updated the stack and want us to remove things.
all_window_ids = []
all_window_targets = []
all_window_tasks = []
if (params.get("DISABLE_PATCHMGMT", "false")).lower() in "true" and params[
"action"
] == "Update":
# they updated the stack and want us to remove things.
patchmgmt.cleanup_patchmgmt(params, boto3_config)

else:
for account_id in account_ids: # across all accounts they desire
window_ids = create_maintenance_window(params, account_id, regions)
window_target_response = define_maintenance_window_targets(params, window_ids["window_ids"], account_id)
window_task_response = define_maintenance_window_tasks(params, window_ids["window_ids"], window_target_response, account_id)
window_id = create_maintenance_window(params, account_id, regions)
all_window_ids.append(window_id["window_ids"])
window_target_response = define_maintenance_window_targets(
params, window_id["window_ids"], account_id
)
all_window_targets.append(window_target_response)
all_window_tasks.append(define_maintenance_window_tasks(
params, window_id["window_ids"], window_target_response, account_id
))
return {
"window_ids": window_ids,
"window_targets": window_target_response,
"window_tasks": window_task_response,
"window_ids": all_window_ids,
"window_targets": all_window_targets,
"window_tasks": all_window_tasks,
}


@helper.create
@helper.update
@helper.delete
def process_cloudformation_event(event: CloudFormationCustomResourceEvent, context: Context) -> str:
def process_cloudformation_event(
event: CloudFormationCustomResourceEvent, context: Context
) -> str:
"""Process Event from AWS CloudFormation.
Args:
Expand All @@ -405,10 +450,9 @@ def process_cloudformation_event(event: CloudFormationCustomResourceEvent, conte
)

if params["action"] in "Add, Update":
response = process_create_update_event(params, regions)
common.store_window_information(response)
process_create_update_event(params, regions)
elif params["action"] == "Remove":
patchmgmt.cleanup_patchmgmt(LOGGER, params, boto3_config)
patchmgmt.cleanup_patchmgmt(params, boto3_config)

return f"sra-patch_mgmt-{params['DELEGATED_ADMIN_ACCOUNT_ID']}"

Expand All @@ -433,4 +477,6 @@ def lambda_handler(event: Dict[str, Any], context: Context) -> None:
helper(event, context)
except Exception:
LOGGER.exception(UNEXPECTED)
raise ValueError(f"Unexpected error executing Lambda function. Review CloudWatch logs '{context.log_group_name}' for details.") from None
raise ValueError(
f"Unexpected error executing Lambda function. Review CloudWatch logs '{context.log_group_name}' for details."
) from None
31 changes: 0 additions & 31 deletions aws_sra_examples/solutions/patch_mgmt/lambda/src/common.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -132,37 +132,6 @@ def get_control_tower_regions() -> list: # noqa: CCR001
return list(customer_regions)


def get_window_information() -> dict: # noqa: CCR001
"""Query SSM Parameter Store to get windows created previously..
Returns:
Window Information Created by this Function to usually be deleted so we don't delete other windows made manually.
"""
ssm_response = SSM_CLIENT.get_parameter(Name="/sra/patch_mgmt/windowInformation")
window_information = ssm_response["Parameter"]["Value"]
return json.loads(window_information)


def store_window_information(window_information: dict) -> bool: # noqa: CCR001
"""Store Window Information for later reference in case of update or delete.
Args:
window_information (dict): Windows that were Created to be stored.
Returns:
Boolean of success or failure
"""
response = SSM_CLIENT.put_parameter(
Name="/sra/patch_mgmt/windowInformation",
Value=json.dumps(window_information),
Description="Created by Patch_Mgmt SRA Solution.",
Type="String",
Overwrite=True,
)
LOGGER.debug({"API_Call": "ssm:PutParameter", "API_Response": response})
return True


def get_enabled_regions(customer_regions: str, control_tower_regions_only: bool = False) -> list: # noqa: CCR001
"""Query STS to identify enabled regions.
Expand Down
81 changes: 45 additions & 36 deletions aws_sra_examples/solutions/patch_mgmt/lambda/src/patchmgmt.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,40 +30,49 @@ def cleanup_patchmgmt(params: dict, boto3_config: Config) -> bool:
Returns:
Boolean of success or failure
"""
window_information = common.get_window_information()
# use boto3 and assume the role to delete all the tasks inside of maintenance windows, then delete the targets, then delete the windows
for window_task in window_information["window_tasks"]:
session = common.assume_role(
params.get("ROLE_NAME_TO_ASSUME", "sra-patch-mgmt-configuration"),
"sra-patch-mgmt-cleanup",
window_task["account_id"],
)
LOGGER.info(f"Deleting Maintenance Window Tasks in {window_task['region']}")
LOGGER.info(window_task)
ssmclient = session.client("ssm", region_name=window_task["region"], config=boto3_config)
response = ssmclient.deregister_task_from_maintenance_window(WindowId=window_task["windowId"], WindowTaskId=window_task["windowTaskId"])
LOGGER.info(response)
for window_target in window_information["window_targets"]:
session = common.assume_role(
params.get("ROLE_NAME_TO_ASSUME", "sra-patch-mgmt-configuration"),
"sra-patch-mgmt-cleanup",
window_target["account_id"],
)
LOGGER.info(f"Deleting Maintenance Window Targets in {window_target['region']}")
LOGGER.info(window_target)
ssmclient = session.client("ssm", region_name=window_target["region"], config=boto3_config)
response = ssmclient.deregister_target_from_maintenance_window(
WindowId=window_target["windowId"],
WindowTargetId=window_target["WindowTargetId"],
)
for previous_window_id in window_information["window_ids"]["windowIds"]:
session = common.assume_role(
params.get("ROLE_NAME_TO_ASSUME", "sra-patch-mgmt-configuration"),
"sra-patch-mgmt-cleanup",
previous_window_id["account_id"],
)
LOGGER.info(f"Deleting Maintenance Windows in {previous_window_id['region']}")
LOGGER.info(previous_window_id)
ssmclient = session.client("ssm", region_name=previous_window_id["region"], config=boto3_config)
response = ssmclient.delete_maintenance_window(WindowId=previous_window_id["windowId"])
account_ids = common.get_account_ids(
[], params["DELEGATED_ADMIN_ACCOUNT_ID"]
)
regions = common.get_enabled_regions(
params.get("ENABLED_REGIONS", ""),
(params.get("CONTROL_TOWER_REGIONS_ONLY", "false")).lower() in "true",
)
for region in regions:
for account in account_ids:
session = common.assume_role(
params.get("ROLE_NAME_TO_ASSUME", "sra-patch-mgmt-configuration"),
"sra-patch-mgmt-cleanup",
account,
)
LOGGER.info(f"Deleting Maintenance Windows in {region}")
ssmclient = session.client("ssm", region_name=region, config=boto3_config)
#Get all maintenance windows, use the next page token if needed, and then loop over each window and check if it has a tag 'createdBy' with a value of 'SRA_Patch_Management' and delete it if so
response = ssmclient.describe_maintenance_windows()
for window in response["WindowIdentities"]:
response2 = ssmclient.list_tags_for_resource(
ResourceType="MaintenanceWindow",
ResourceId=window["WindowId"])
#For tag in tag list then check if the tag is 'createdBy' and if it is then delete the window
for tag in response2["TagList"]:
if tag["Key"] == "createdBy" and tag["Value"] == "SRA_Patch_Management":
ssmclient.delete_maintenance_window(WindowId=window["WindowId"])
LOGGER.info(f"Deleted Maintenance Window {window['Name']}")
break
while "NextToken" in response:
response = ssmclient.describe_maintenance_windows(
NextToken=response["NextToken"]
)
for window in response["WindowIdentities"]:
#Use the Window ID to list_tags_for_resource and then check if it has a tag 'createdBy' with a value of 'SRA_Patch_Management' and delete it if so
response2 = ssmclient.list_tags_for_resource(
ResourceType="MaintenanceWindow",
ResourceId=window["WindowId"])
#For tag in tag list then check if the tag is 'createdBy' and if it is then delete the window
for tag in response2["TagList"]:
if tag["Key"] == "createdBy" and tag["Value"] == "SRA_Patch_Management":
ssmclient.delete_maintenance_window(WindowId=window["WindowId"])
LOGGER.info(f"Deleted Maintenance Window {window['Name']}")
break


return True
Loading

0 comments on commit a67f6c5

Please sign in to comment.