-
Notifications
You must be signed in to change notification settings - Fork 329
Expand file tree
/
Copy pathlambda_ec2_stop_cron.py
More file actions
180 lines (147 loc) · 4.96 KB
/
lambda_ec2_stop_cron.py
File metadata and controls
180 lines (147 loc) · 4.96 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
"""
-*- coding: utf-8 -*-
========================
AWS Lambda
========================
Contributor: Chirag Rathod (Srce Cde)
========================
"""
import os
import sys
import traceback
import logging
import ast
import json
import boto3
logger = logging.getLogger()
logger.setLevel(logging.INFO)
# retrieve list of instance ids from ENV variable to skip stopping
INSTANCE_IDS_TO_IGNORE_STOP = ast.literal_eval(
os.environ.get("INSTANCE_IDS_TO_IGNORE_STOP", "[]")
)
def process_error() -> dict:
ex_type, ex_value, ex_traceback = sys.exc_info()
traceback_string = traceback.format_exception(ex_type, ex_value, ex_traceback)
error_msg = json.dumps(
{
"errorType": ex_type.__name__,
"errorMessage": str(ex_value),
"stackTrace": traceback_string,
}
)
return error_msg
def fetch_regions() -> list:
"""
Helper function to retrieve regions
Returns:
--------
regions: list of AWS regions
"""
ec2_client = boto3.client("ec2")
try:
regions = ec2_client.describe_regions()
except:
error_msg = process_error()
logger.error(error_msg)
return regions["Regions"]
def stop_instances(ec2_client: object, instance_ids: list) -> None:
"""
Helper function to stop the instances
Parameters:
-----------
ec2_client: boto3 region specific object
instance_ids: list of instance ids to stop
"""
try:
response = ec2_client.stop_instances(InstanceIds=instance_ids, Force=True)
except:
error_msg = process_error()
logger.error(error_msg)
def start_instances(ec2_client: object, instance_ids: list) -> None:
"""
Helper function to stop the instances
Parameters:
-----------
ec2_client: boto3 region specific object
instance_ids: list of instance ids to start
"""
try:
response = ec2_client.start_instances(InstanceIds=instance_ids)
except:
error_msg = process_error()
logger.error(error_msg)
def get_instance_ids(response: dict, STOP: bool) -> list:
"""
Parse the instance IDs from response
Parameters:
-----------
response: boto3 describe_instances response
STOP: Flag to decide type of instance Ids to fetch
Returns:
--------
instance_ids: list of instance ids to stop
"""
if STOP:
instance_ids_to_stop = []
for resp in response.get("Reservations"):
for instance in resp.get("Instances"):
if (
instance.get("State").get("Name") in ["pending", "running"]
and instance["InstanceId"] not in INSTANCE_IDS_TO_IGNORE_STOP
):
instance_ids_to_stop.append(instance["InstanceId"])
return instance_ids_to_stop
else:
instance_ids_to_start = []
for resp in response.get("Reservations"):
for instance in resp.get("Instances"):
if instance.get("State").get("Name") in ["stopped", "stopping"]:
instance_ids_to_start.append(instance["InstanceId"])
return instance_ids_to_start
def start_stop_instances_across_region(regions: list, STOP=True) -> None:
"""
Start / Stop the instances across regions
Parameters:
-----------
regions: list of regions to analyze
STOP: Flag to decide whether to start or stop the instances
"""
try:
for region in regions:
ec2_client = boto3.client("ec2", region_name=region["RegionName"])
response = ec2_client.describe_instances()
instance_ids = get_instance_ids(response, STOP)
if instance_ids and STOP:
stop_instances(ec2_client, instance_ids)
if instance_ids and not STOP:
start_instances(ec2_client, instance_ids)
while "NextToken" in response:
response = ec2_client.describe_instances(
NextToken=response["NextToken"]
)
instance_ids = get_instance_ids(response, STOP)
if instance_ids and STOP:
stop_instances(ec2_client, instance_ids)
if instance_ids and not STOP:
start_instances(ec2_client, instance_ids)
except:
error_msg = process_error()
logger.error(error_msg)
def lambda_handler(event, context):
"""
Main handler
"""
logging.info(event)
# retrieve regions
regions = fetch_regions()
try:
rule_type = event["resources"][0].split("/")[-1]
# checking which rule triggered the lambda function
if "ScheduledEC2StopRule" in rule_type:
start_stop_instances_across_region(regions, STOP=True)
if "ScheduledEC2StartRule" in rule_type:
start_stop_instances_across_region(regions, STOP=False)
except:
error_msg = process_error()
logger.error(error_msg)
return {"statusCode": 200, "body": json.dumps("Thanks from Srce Cde (Chirag)!")}