Skip to content

Commit ac06368

Browse files
tmekaritmekari
authored
Adding all endpoint types to SNS.2 Remediation (#145) (#147)
* Adding all endpoint types * Fixed test, added error logging to reset
1 parent 001ec08 commit ac06368

File tree

3 files changed

+111
-55
lines changed

3 files changed

+111
-55
lines changed

‎source/remediation_runbooks/scripts/enable_delivery_status_logging.py‎

Lines changed: 30 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,7 @@
1111
}
1212
)
1313

14-
failureFeedbackRoleValue = "LambdaFailureFeedbackRoleArn"
15-
successFeedbackRoleValue = "LambdaSuccessFeedbackRoleArn"
16-
successRateRoleValue = "LambdaSuccessFeedbackSampleRate"
14+
endpointTypes = ['HTTP', 'Firehose', 'Lambda', 'Application', 'SQS']
1715

1816
def connect_to_sns():
1917
return boto3.client('sns', config=boto_config)
@@ -39,9 +37,21 @@ def lambda_handler(event, _):
3937
topic_attributes = get_topic_attributes(topic_arn)
4038

4139
return {
42-
"FailureFeedbackRole": topic_attributes["Attributes"][failureFeedbackRoleValue],
43-
"SuccessFeedbackRole": topic_attributes["Attributes"][successFeedbackRoleValue],
44-
"SuccessSampleRate": topic_attributes["Attributes"][successRateRoleValue]
40+
"HTTPFailureFeedbackRoleArn": topic_attributes["Attributes"]["HTTPFailureFeedbackRoleArn"],
41+
"HTTPSuccessFeedbackRoleArn": topic_attributes["Attributes"]["HTTPSuccessFeedbackRoleArn"],
42+
"HTTPSuccessFeedbackSampleRate": topic_attributes["Attributes"]["HTTPSuccessFeedbackSampleRate"],
43+
"FirehoseFailureFeedbackRoleArn": topic_attributes["Attributes"]["FirehoseFailureFeedbackRoleArn"],
44+
"FirehoseSuccessFeedbackRoleArn": topic_attributes["Attributes"]["FirehoseSuccessFeedbackRoleArn"],
45+
"FirehoseSuccessFeedbackSampleRate": topic_attributes["Attributes"]["FirehoseSuccessFeedbackSampleRate"],
46+
"LambdaFailureFeedbackRoleArn": topic_attributes["Attributes"]["LambdaFailureFeedbackRoleArn"],
47+
"LambdaSuccessFeedbackRoleArn": topic_attributes["Attributes"]["LambdaSuccessFeedbackRoleArn"],
48+
"LambdaSuccessFeedbackSampleRate": topic_attributes["Attributes"]["LambdaSuccessFeedbackSampleRate"],
49+
"ApplicationFailureFeedbackRoleArn": topic_attributes["Attributes"]["ApplicationFailureFeedbackRoleArn"],
50+
"ApplicationSuccessFeedbackRoleArn": topic_attributes["Attributes"]["ApplicationSuccessFeedbackRoleArn"],
51+
"ApplicationSuccessFeedbackSampleRate": topic_attributes["Attributes"]["ApplicationSuccessFeedbackSampleRate"],
52+
"SQSFailureFeedbackRoleArn": topic_attributes["Attributes"]["SQSFailureFeedbackRoleArn"],
53+
"SQSSuccessFeedbackRoleArn": topic_attributes["Attributes"]["SQSSuccessFeedbackRoleArn"],
54+
"SQSSuccessFeedbackSampleRate": topic_attributes["Attributes"]["SQSSuccessFeedbackSampleRate"]
4555
}
4656

4757
def add_roles_to_topic(logging_role, topic_arn):
@@ -50,24 +60,26 @@ def add_roles_to_topic(logging_role, topic_arn):
5060
"""
5161
sns = connect_to_sns()
5262
try:
53-
sns.set_topic_attributes(TopicArn=topic_arn, AttributeName=successFeedbackRoleValue, AttributeValue=logging_role)
54-
sns.set_topic_attributes(TopicArn=topic_arn, AttributeName=failureFeedbackRoleValue, AttributeValue=logging_role)
63+
for endpoint in endpointTypes:
64+
sns.set_topic_attributes(TopicArn=topic_arn, AttributeName=f'{endpoint}SuccessFeedbackRoleArn', AttributeValue=logging_role)
65+
sns.set_topic_attributes(TopicArn=topic_arn, AttributeName=f'{endpoint}FailureFeedbackRoleArn', AttributeValue=logging_role)
5566

5667
except Exception as e:
5768
reset_to_recognized_state(topic_arn)
58-
exit(f'Failed to set success/failure role of topic '+topic_arn+': '+str(e))
69+
exit(f'Failed to set success/failure role of topic {topic_arn}: {str(e)}')
5970

6071
def add_sample_rate_to_topic(topic_arn, sample_rate):
6172
"""
6273
Configures the Success sample rate, the percentage of successful messages for which you want to receive CloudWatch Logs.
6374
"""
6475
sns = connect_to_sns()
6576
try:
66-
sns.set_topic_attributes(TopicArn=topic_arn, AttributeName=successRateRoleValue, AttributeValue=sample_rate)
77+
for endpoint in endpointTypes:
78+
sns.set_topic_attributes(TopicArn=topic_arn, AttributeName=f'{endpoint}SuccessFeedbackSampleRate', AttributeValue=sample_rate)
6779

6880
except Exception as e:
6981
reset_to_recognized_state(topic_arn)
70-
exit(f'Failed to set success sample rate of SNS topic '+topic_arn+': '+str(e))
82+
exit(f'Failed to set success sample rate of SNS topic {topic_arn}: {str(e)}')
7183

7284
def get_topic_attributes(topic_arn):
7385
"""
@@ -79,13 +91,16 @@ def get_topic_attributes(topic_arn):
7991
return topic_attributes
8092

8193
except Exception as e:
82-
exit(f'Failed to get attributes of SNS topic '+topic_arn+': '+str(e))
94+
exit(f'Failed to get attributes of SNS topic {topic_arn}: {str(e)}')
8395

8496
def reset_to_recognized_state(topic_arn):
8597
"""
8698
Used in case of error, will unset all delivery status logging parameters.
8799
"""
88100
sns = connect_to_sns()
89-
90-
sns.set_topic_attributes(TopicArn=topic_arn, AttributeName=successFeedbackRoleValue, AttributeValue='')
91-
sns.set_topic_attributes(TopicArn=topic_arn, AttributeName=failureFeedbackRoleValue, AttributeValue='')
101+
for endpoint in endpointTypes:
102+
try:
103+
sns.set_topic_attributes(TopicArn=topic_arn, AttributeName=f'{endpoint}SuccessFeedbackRoleArn', AttributeValue='')
104+
sns.set_topic_attributes(TopicArn=topic_arn, AttributeName=f'{endpoint}FailureFeedbackRoleArn', AttributeValue='')
105+
except Exception:
106+
print(f'There was an error while resetting SNS Topic {topic_arn}, please manually turn off delivery status logging for protocol {endpoint}')

‎source/remediation_runbooks/scripts/test/test_enable_delivery_status_logging.py‎

Lines changed: 49 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
from enable_delivery_status_logging import lambda_handler
1010

1111
def test_enables_delivery_status_logging(mocker):
12+
endpointTypes = ['HTTP', 'Firehose', 'Lambda', 'Application', 'SQS']
1213

1314
my_session = boto3.session.Session()
1415
my_region = my_session.region_name
@@ -26,34 +27,47 @@ def test_enables_delivery_status_logging(mocker):
2627
logging_arn = 'logging_arn'
2728
topic_arn = f'arn:aws:sns:{my_region}:111111111111:sharr-test'
2829

29-
response = { 'Attributes' : {
30+
stub_response = { 'Attributes' : {
3031
"LambdaFailureFeedbackRoleArn": logging_arn,
3132
"LambdaSuccessFeedbackRoleArn": logging_arn,
32-
"LambdaSuccessFeedbackSampleRate": logging_arn
33+
"LambdaSuccessFeedbackSampleRate": '0',
34+
"HTTPFailureFeedbackRoleArn": logging_arn,
35+
"HTTPSuccessFeedbackRoleArn": logging_arn,
36+
"HTTPSuccessFeedbackSampleRate": '0',
37+
"FirehoseFailureFeedbackRoleArn": logging_arn,
38+
"FirehoseSuccessFeedbackRoleArn": logging_arn,
39+
"FirehoseSuccessFeedbackSampleRate": '0',
40+
"ApplicationFailureFeedbackRoleArn": logging_arn,
41+
"ApplicationSuccessFeedbackRoleArn": logging_arn,
42+
"ApplicationSuccessFeedbackSampleRate": '0',
43+
"SQSFailureFeedbackRoleArn": logging_arn,
44+
"SQSSuccessFeedbackRoleArn": logging_arn,
45+
"SQSSuccessFeedbackSampleRate": '0'
3346
}}
47+
for endpoint in endpointTypes:
48+
stub_sns.add_response(
49+
'set_topic_attributes',
50+
{},
51+
{
52+
'TopicArn': topic_arn ,
53+
'AttributeName': f"{endpoint}SuccessFeedbackRoleArn",
54+
'AttributeValue': logging_arn})
3455

35-
stub_sns.add_response(
36-
'set_topic_attributes',
56+
stub_sns.add_response('set_topic_attributes',
3757
{},
38-
{
39-
'TopicArn': topic_arn ,
40-
'AttributeName': "LambdaSuccessFeedbackRoleArn",
58+
{ 'TopicArn': topic_arn ,
59+
'AttributeName': f"{endpoint}FailureFeedbackRoleArn" ,
4160
'AttributeValue': logging_arn})
42-
43-
stub_sns.add_response('set_topic_attributes',
44-
{},
45-
{ 'TopicArn': topic_arn ,
46-
'AttributeName': "LambdaFailureFeedbackRoleArn" ,
47-
'AttributeValue': logging_arn})
48-
49-
stub_sns.add_response('set_topic_attributes',
50-
{},
51-
{ 'TopicArn': topic_arn ,
52-
'AttributeName': "LambdaSuccessFeedbackSampleRate" ,
53-
'AttributeValue': '0'})
61+
62+
for endpoint in endpointTypes:
63+
stub_sns.add_response('set_topic_attributes',
64+
{},
65+
{ 'TopicArn': topic_arn ,
66+
'AttributeName': f"{endpoint}SuccessFeedbackSampleRate" ,
67+
'AttributeValue': '0'})
5468

5569
stub_sns.add_response('get_topic_attributes',
56-
response,
70+
stub_response,
5771
{ 'TopicArn': topic_arn })
5872

5973
stub_sns.activate()
@@ -62,8 +76,20 @@ def test_enables_delivery_status_logging(mocker):
6276
event = { 'topic_arn': topic_arn, 'logging_role': logging_arn, 'sample_rate': '0' }
6377
response = lambda_handler(event, {})
6478
assert response == {
65-
"FailureFeedbackRole": logging_arn,
66-
"SuccessFeedbackRole": logging_arn,
67-
"SuccessSampleRate": logging_arn }
79+
"LambdaFailureFeedbackRoleArn": logging_arn,
80+
"LambdaSuccessFeedbackRoleArn": logging_arn,
81+
"LambdaSuccessFeedbackSampleRate": '0',
82+
"HTTPFailureFeedbackRoleArn": logging_arn,
83+
"HTTPSuccessFeedbackRoleArn": logging_arn,
84+
"HTTPSuccessFeedbackSampleRate": '0',
85+
"FirehoseFailureFeedbackRoleArn": logging_arn,
86+
"FirehoseSuccessFeedbackRoleArn": logging_arn,
87+
"FirehoseSuccessFeedbackSampleRate": '0',
88+
"ApplicationFailureFeedbackRoleArn": logging_arn,
89+
"ApplicationSuccessFeedbackRoleArn": logging_arn,
90+
"ApplicationSuccessFeedbackSampleRate": '0',
91+
"SQSFailureFeedbackRoleArn": logging_arn,
92+
"SQSSuccessFeedbackRoleArn": logging_arn,
93+
"SQSSuccessFeedbackSampleRate": '0' }
6894

6995

‎source/test/__snapshots__/runbook_stack.test.ts.snap‎

Lines changed: 32 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -3872,9 +3872,7 @@ boto_config = Config(
38723872
}
38733873
)
38743874

3875-
failureFeedbackRoleValue = "LambdaFailureFeedbackRoleArn"
3876-
successFeedbackRoleValue = "LambdaSuccessFeedbackRoleArn"
3877-
successRateRoleValue = "LambdaSuccessFeedbackSampleRate"
3875+
endpointTypes = ['HTTP', 'Firehose', 'Lambda', 'Application', 'SQS']
38783876

38793877
def connect_to_sns():
38803878
return boto3.client('sns', config=boto_config)
@@ -3900,9 +3898,21 @@ def lambda_handler(event, _):
39003898
topic_attributes = get_topic_attributes(topic_arn)
39013899

39023900
return {
3903-
"FailureFeedbackRole": topic_attributes["Attributes"][failureFeedbackRoleValue],
3904-
"SuccessFeedbackRole": topic_attributes["Attributes"][successFeedbackRoleValue],
3905-
"SuccessSampleRate": topic_attributes["Attributes"][successRateRoleValue]
3901+
"HTTPFailureFeedbackRoleArn": topic_attributes["Attributes"]["HTTPFailureFeedbackRoleArn"],
3902+
"HTTPSuccessFeedbackRoleArn": topic_attributes["Attributes"]["HTTPSuccessFeedbackRoleArn"],
3903+
"HTTPSuccessFeedbackSampleRate": topic_attributes["Attributes"]["HTTPSuccessFeedbackSampleRate"],
3904+
"FirehoseFailureFeedbackRoleArn": topic_attributes["Attributes"]["FirehoseFailureFeedbackRoleArn"],
3905+
"FirehoseSuccessFeedbackRoleArn": topic_attributes["Attributes"]["FirehoseSuccessFeedbackRoleArn"],
3906+
"FirehoseSuccessFeedbackSampleRate": topic_attributes["Attributes"]["FirehoseSuccessFeedbackSampleRate"],
3907+
"LambdaFailureFeedbackRoleArn": topic_attributes["Attributes"]["LambdaFailureFeedbackRoleArn"],
3908+
"LambdaSuccessFeedbackRoleArn": topic_attributes["Attributes"]["LambdaSuccessFeedbackRoleArn"],
3909+
"LambdaSuccessFeedbackSampleRate": topic_attributes["Attributes"]["LambdaSuccessFeedbackSampleRate"],
3910+
"ApplicationFailureFeedbackRoleArn": topic_attributes["Attributes"]["ApplicationFailureFeedbackRoleArn"],
3911+
"ApplicationSuccessFeedbackRoleArn": topic_attributes["Attributes"]["ApplicationSuccessFeedbackRoleArn"],
3912+
"ApplicationSuccessFeedbackSampleRate": topic_attributes["Attributes"]["ApplicationSuccessFeedbackSampleRate"],
3913+
"SQSFailureFeedbackRoleArn": topic_attributes["Attributes"]["SQSFailureFeedbackRoleArn"],
3914+
"SQSSuccessFeedbackRoleArn": topic_attributes["Attributes"]["SQSSuccessFeedbackRoleArn"],
3915+
"SQSSuccessFeedbackSampleRate": topic_attributes["Attributes"]["SQSSuccessFeedbackSampleRate"]
39063916
}
39073917

39083918
def add_roles_to_topic(logging_role, topic_arn):
@@ -3911,24 +3921,26 @@ def add_roles_to_topic(logging_role, topic_arn):
39113921
"""
39123922
sns = connect_to_sns()
39133923
try:
3914-
sns.set_topic_attributes(TopicArn=topic_arn, AttributeName=successFeedbackRoleValue, AttributeValue=logging_role)
3915-
sns.set_topic_attributes(TopicArn=topic_arn, AttributeName=failureFeedbackRoleValue, AttributeValue=logging_role)
3924+
for endpoint in endpointTypes:
3925+
sns.set_topic_attributes(TopicArn=topic_arn, AttributeName=f'{endpoint}SuccessFeedbackRoleArn', AttributeValue=logging_role)
3926+
sns.set_topic_attributes(TopicArn=topic_arn, AttributeName=f'{endpoint}FailureFeedbackRoleArn', AttributeValue=logging_role)
39163927

39173928
except Exception as e:
39183929
reset_to_recognized_state(topic_arn)
3919-
exit(f'Failed to set success/failure role of topic '+topic_arn+': '+str(e))
3930+
exit(f'Failed to set success/failure role of topic {topic_arn}: {str(e)}')
39203931

39213932
def add_sample_rate_to_topic(topic_arn, sample_rate):
39223933
"""
39233934
Configures the Success sample rate, the percentage of successful messages for which you want to receive CloudWatch Logs.
39243935
"""
39253936
sns = connect_to_sns()
39263937
try:
3927-
sns.set_topic_attributes(TopicArn=topic_arn, AttributeName=successRateRoleValue, AttributeValue=sample_rate)
3938+
for endpoint in endpointTypes:
3939+
sns.set_topic_attributes(TopicArn=topic_arn, AttributeName=f'{endpoint}SuccessFeedbackSampleRate', AttributeValue=sample_rate)
39283940

39293941
except Exception as e:
39303942
reset_to_recognized_state(topic_arn)
3931-
exit(f'Failed to set success sample rate of SNS topic '+topic_arn+': '+str(e))
3943+
exit(f'Failed to set success sample rate of SNS topic {topic_arn}: {str(e)}')
39323944

39333945
def get_topic_attributes(topic_arn):
39343946
"""
@@ -3940,16 +3952,19 @@ def get_topic_attributes(topic_arn):
39403952
return topic_attributes
39413953

39423954
except Exception as e:
3943-
exit(f'Failed to get attributes of SNS topic '+topic_arn+': '+str(e))
3955+
exit(f'Failed to get attributes of SNS topic {topic_arn}: {str(e)}')
39443956

39453957
def reset_to_recognized_state(topic_arn):
39463958
"""
39473959
Used in case of error, will unset all delivery status logging parameters.
39483960
"""
39493961
sns = connect_to_sns()
3950-
3951-
sns.set_topic_attributes(TopicArn=topic_arn, AttributeName=successFeedbackRoleValue, AttributeValue='')
3952-
sns.set_topic_attributes(TopicArn=topic_arn, AttributeName=failureFeedbackRoleValue, AttributeValue='')",
3962+
for endpoint in endpointTypes:
3963+
try:
3964+
sns.set_topic_attributes(TopicArn=topic_arn, AttributeName=f'{endpoint}SuccessFeedbackRoleArn', AttributeValue='')
3965+
sns.set_topic_attributes(TopicArn=topic_arn, AttributeName=f'{endpoint}FailureFeedbackRoleArn', AttributeValue='')
3966+
except Exception:
3967+
print(f'There was an error while resetting SNS Topic {topic_arn}, please manually turn off delivery status logging for protocol {endpoint}')",
39533968
},
39543969
"isEnd": true,
39553970
"name": "EnableDeliveryStatusLogging",
@@ -7737,7 +7752,7 @@ def add_ssl_bucket_policy(event, _):
77377752
"Properties": {
77387753
"CreateIntervalSeconds": 1,
77397754
"DeleteIntervalSeconds": 0,
7740-
"DocumentPropertiesHash": "76a143310018a965418e9379eb90670ea780456ea33047bd55f0b597443d59ba",
7755+
"DocumentPropertiesHash": "dff26ffe5bc8a2eb7a84297f591aa46181b17efa8b7c977636db9cf3beeec61c",
77417756
"ServiceToken": {
77427757
"Ref": "WaitProviderServiceToken",
77437758
},
@@ -7931,7 +7946,7 @@ def add_ssl_bucket_policy(event, _):
79317946
"Properties": {
79327947
"CreateIntervalSeconds": 0,
79337948
"DeleteIntervalSeconds": 0.5,
7934-
"DocumentPropertiesHash": "76a143310018a965418e9379eb90670ea780456ea33047bd55f0b597443d59ba",
7949+
"DocumentPropertiesHash": "dff26ffe5bc8a2eb7a84297f591aa46181b17efa8b7c977636db9cf3beeec61c",
79357950
"ServiceToken": {
79367951
"Ref": "WaitProviderServiceToken",
79377952
},

0 commit comments

Comments
 (0)