44import json
55import time
66import unittest
7+ import copy
78from io import StringIO
89
9- from firetail_lambda import firetail_handler
10+ from firetail_lambda import firetail_handler, firetail_app
1011
12+ api_gateway_v1 = {
13+ "body": "eyJ0ZXN0IjoiYm9keSJ9",
14+ "resource": "/{proxy+}",
15+ "path": "/path/to/resource",
16+ "httpMethod": "POST",
17+ "isBase64Encoded": True,
18+ "queryStringParameters": {
19+ "foo": "bar"
20+ },
21+ "multiValueQueryStringParameters": {
22+ "foo": [
23+ "bar"
24+ ]
25+ },
26+ "pathParameters": {
27+ "proxy": "/path/to/resource"
28+ },
29+ "stageVariables": {
30+ "baz": "qux"
31+ },
32+ "headers": {
33+ "Authorization": "Bearer jwt123.123.123",
34+ "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8",
35+ "Accept-Encoding": "gzip, deflate, sdch",
36+ "Accept-Language": "en-US,en;q=0.8",
37+ "Cache-Control": "max-age=0",
38+ "CloudFront-Forwarded-Proto": "https",
39+ "CloudFront-Is-Desktop-Viewer": "true",
40+ "CloudFront-Is-Mobile-Viewer": "false",
41+ "CloudFront-Is-SmartTV-Viewer": "false",
42+ "CloudFront-Is-Tablet-Viewer": "false",
43+ "CloudFront-Viewer-Country": "US",
44+ "Host": "1234567890.execute-api.us-east-1.amazonaws.com",
45+ "Upgrade-Insecure-Requests": "1",
46+ "User-Agent": "Custom User Agent String",
47+ "Via": "1.1 08f323deadbeefa7af34d5feb414ce27.cloudfront.net (CloudFront)",
48+ "X-Amz-Cf-Id": "cDehVQoZnx43VYQb9j2-nvCh-9z396Uhbp027Y2JvkCPNLmGJHqlaA==",
49+ "X-Forwarded-For": "127.0.0.1, 127.0.0.2",
50+ "X-Forwarded-Port": "443",
51+ "X-Forwarded-Proto": "https"
52+ },
53+ "multiValueHeaders": {
54+ "Authorization":[
55+ "Bearer jwt123.123.123"
56+ ],
57+ "Accept": [
58+ "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8"
59+ ],
60+ "Accept-Encoding": [
61+ "gzip, deflate, sdch"
62+ ],
63+ "Accept-Language": [
64+ "en-US,en;q=0.8"
65+ ],
66+ "Cache-Control": [
67+ "max-age=0"
68+ ],
69+ "CloudFront-Forwarded-Proto": [
70+ "https"
71+ ],
72+ "CloudFront-Is-Desktop-Viewer": [
73+ "true"
74+ ],
75+ "CloudFront-Is-Mobile-Viewer": [
76+ "false"
77+ ],
78+ "CloudFront-Is-SmartTV-Viewer": [
79+ "false"
80+ ],
81+ "CloudFront-Is-Tablet-Viewer": [
82+ "false"
83+ ],
84+ "CloudFront-Viewer-Country": [
85+ "US"
86+ ],
87+ "Host": [
88+ "0123456789.execute-api.us-east-1.amazonaws.com"
89+ ],
90+ "Upgrade-Insecure-Requests": [
91+ "1"
92+ ],
93+ "User-Agent": [
94+ "Custom User Agent String"
95+ ],
96+ "Via": [
97+ "1.1 08f323deadbeefa7af34d5feb414ce27.cloudfront.net (CloudFront)"
98+ ],
99+ "X-Amz-Cf-Id": [
100+ "cDehVQoZnx43VYQb9j2-nvCh-9z396Uhbp027Y2JvkCPNLmGJHqlaA=="
101+ ],
102+ "X-Forwarded-For": [
103+ "127.0.0.1, 127.0.0.2"
104+ ],
105+ "X-Forwarded-Port": [
106+ "443"
107+ ],
108+ "X-Forwarded-Proto": [
109+ "https"
110+ ]
111+ },
112+ "requestContext": {
113+ "accountId": "123456789012",
114+ "resourceId": "123456",
115+ "stage": "prod",
116+ "requestId": "c6af9ac6-7b61-11e6-9a41-93e8deadbeef",
117+ "requestTime": "09/Apr/2015:12:34:56 +0000",
118+ "requestTimeEpoch": 1428582896000,
119+ "identity": {
120+ "cognitoIdentityPoolId": None,
121+ "accountId": None,
122+ "cognitoIdentityId": None,
123+ "caller": None,
124+ "accessKey": None,
125+ "sourceIp": "127.0.0.1",
126+ "cognitoAuthenticationType": None,
127+ "cognitoAuthenticationProvider": None,
128+ "userArn": None,
129+ "userAgent": "Custom User Agent String",
130+ "user": None
131+ },
132+ "path": "/prod/path/to/resource",
133+ "resourcePath": "/{proxy+}",
134+ "httpMethod": "POST",
135+ "apiId": "1234567890",
136+ "protocol": "HTTP/1.1"
137+ }
138+ }
11139
12140class TestSimple(unittest.TestCase):
13141
14142 def test_handler_api(self):
15143 event = {}
16- @firetail_handler()
144+ app = firetail_app()
145+ @firetail_handler(app)
17146 def handler(event, context):
18147 return 201, json.dumps({"message": "success"})
19148
@@ -25,7 +154,8 @@ def handler(event, context):
25154
26155 def test_incorrect_handler_api(self):
27156 event = {}
28- @firetail_handler()
157+ app = firetail_app()
158+ @firetail_handler(app)
29159 def handler(argument):
30160 return 201, json.dumps({"message": "success"})
31161
@@ -37,7 +167,9 @@ def handler(argument):
37167
38168 def test_handler_sleeper_api(self):
39169 event = {}
40- @firetail_handler(enable_sleeper=True)
170+ app = firetail_app()
171+ app.enable_sleeper = True
172+ @firetail_handler(app)
41173 def handler(event, context):
42174 return 201, json.dumps({"message": "success"})
43175
@@ -52,6 +184,32 @@ def handler(event, context):
52184 self.assertGreaterEqual(difference, .5)
53185 self.assertEqual(output, 'firetail:loggingapi:eyJldmVudCI6IHt9LCAicmVzcG9uc2UiOiBbMjAxLCAie1wibWVzc2FnZVwiOiBcInN1Y2Nlc3NcIn0iXX0=')
54186
187+ def test_handler_sanitizer(self):
188+ event = api_gateway_v1
189+ def sanitize_payloads(event, response):
190+ new_event = copy.copy(event)
191+ remove_headers = ['authorization','Authorization', 'x-api-key']
192+ if 'headers' in event:
193+ for header in remove_headers:
194+ if header in event['headers']:
195+ del new_event['headers'][header]
196+ if 'multiValueHeaders' in event and header in event['multiValueHeaders']:
197+ del new_event['multiValueHeaders'][header]
198+
199+ return new_event, response
200+ app = firetail_app()
201+ app.sanitization_callback = sanitize_payloads
202+ @firetail_handler(app)
203+ def handler(event, context):
204+ return 201, json.dumps({"message": "success"})
205+
206+ temp_stdout = StringIO()
207+ with contextlib.redirect_stdout(temp_stdout):
208+ handler(event, "")
209+
210+ output = temp_stdout.getvalue().strip()
211+ self.assertEqual(output, 'firetail:loggingapi:eyJldmVudCI6IHsiYm9keSI6ICJleUowWlhOMElqb2lZbTlrZVNKOSIsICJyZXNvdXJjZSI6ICIve3Byb3h5K30iLCAicGF0aCI6ICIvcGF0aC90by9yZXNvdXJjZSIsICJodHRwTWV0aG9kIjogIlBPU1QiLCAiaXNCYXNlNjRFbmNvZGVkIjogdHJ1ZSwgInF1ZXJ5U3RyaW5nUGFyYW1ldGVycyI6IHsiZm9vIjogImJhciJ9LCAibXVsdGlWYWx1ZVF1ZXJ5U3RyaW5nUGFyYW1ldGVycyI6IHsiZm9vIjogWyJiYXIiXX0sICJwYXRoUGFyYW1ldGVycyI6IHsicHJveHkiOiAiL3BhdGgvdG8vcmVzb3VyY2UifSwgInN0YWdlVmFyaWFibGVzIjogeyJiYXoiOiAicXV4In0sICJoZWFkZXJzIjogeyJBY2NlcHQiOiAidGV4dC9odG1sLGFwcGxpY2F0aW9uL3hodG1sK3htbCxhcHBsaWNhdGlvbi94bWw7cT0wLjksaW1hZ2Uvd2VicCwqLyo7cT0wLjgiLCAiQWNjZXB0LUVuY29kaW5nIjogImd6aXAsIGRlZmxhdGUsIHNkY2giLCAiQWNjZXB0LUxhbmd1YWdlIjogImVuLVVTLGVuO3E9MC44IiwgIkNhY2hlLUNvbnRyb2wiOiAibWF4LWFnZT0wIiwgIkNsb3VkRnJvbnQtRm9yd2FyZGVkLVByb3RvIjogImh0dHBzIiwgIkNsb3VkRnJvbnQtSXMtRGVza3RvcC1WaWV3ZXIiOiAidHJ1ZSIsICJDbG91ZEZyb250LUlzLU1vYmlsZS1WaWV3ZXIiOiAiZmFsc2UiLCAiQ2xvdWRGcm9udC1Jcy1TbWFydFRWLVZpZXdlciI6ICJmYWxzZSIsICJDbG91ZEZyb250LUlzLVRhYmxldC1WaWV3ZXIiOiAiZmFsc2UiLCAiQ2xvdWRGcm9udC1WaWV3ZXItQ291bnRyeSI6ICJVUyIsICJIb3N0IjogIjEyMzQ1Njc4OTAuZXhlY3V0ZS1hcGkudXMtZWFzdC0xLmFtYXpvbmF3cy5jb20iLCAiVXBncmFkZS1JbnNlY3VyZS1SZXF1ZXN0cyI6ICIxIiwgIlVzZXItQWdlbnQiOiAiQ3VzdG9tIFVzZXIgQWdlbnQgU3RyaW5nIiwgIlZpYSI6ICIxLjEgMDhmMzIzZGVhZGJlZWZhN2FmMzRkNWZlYjQxNGNlMjcuY2xvdWRmcm9udC5uZXQgKENsb3VkRnJvbnQpIiwgIlgtQW16LUNmLUlkIjogImNEZWhWUW9abng0M1ZZUWI5ajItbnZDaC05ejM5NlVoYnAwMjdZMkp2a0NQTkxtR0pIcWxhQT09IiwgIlgtRm9yd2FyZGVkLUZvciI6ICIxMjcuMC4wLjEsIDEyNy4wLjAuMiIsICJYLUZvcndhcmRlZC1Qb3J0IjogIjQ0MyIsICJYLUZvcndhcmRlZC1Qcm90byI6ICJodHRwcyJ9LCAibXVsdGlWYWx1ZUhlYWRlcnMiOiB7IkFjY2VwdCI6IFsidGV4dC9odG1sLGFwcGxpY2F0aW9uL3hodG1sK3htbCxhcHBsaWNhdGlvbi94bWw7cT0wLjksaW1hZ2Uvd2VicCwqLyo7cT0wLjgiXSwgIkFjY2VwdC1FbmNvZGluZyI6IFsiZ3ppcCwgZGVmbGF0ZSwgc2RjaCJdLCAiQWNjZXB0LUxhbmd1YWdlIjogWyJlbi1VUyxlbjtxPTAuOCJdLCAiQ2FjaGUtQ29udHJvbCI6IFsibWF4LWFnZT0wIl0sICJDbG91ZEZyb250LUZvcndhcmRlZC1Qcm90byI6IFsiaHR0cHMiXSwgIkNsb3VkRnJvbnQtSXMtRGVza3RvcC1WaWV3ZXIiOiBbInRydWUiXSwgIkNsb3VkRnJvbnQtSXMtTW9iaWxlLVZpZXdlciI6IFsiZmFsc2UiXSwgIkNsb3VkRnJvbnQtSXMtU21hcnRUVi1WaWV3ZXIiOiBbImZhbHNlIl0sICJDbG91ZEZyb250LUlzLVRhYmxldC1WaWV3ZXIiOiBbImZhbHNlIl0sICJDbG91ZEZyb250LVZpZXdlci1Db3VudHJ5IjogWyJVUyJdLCAiSG9zdCI6IFsiMDEyMzQ1Njc4OS5leGVjdXRlLWFwaS51cy1lYXN0LTEuYW1hem9uYXdzLmNvbSJdLCAiVXBncmFkZS1JbnNlY3VyZS1SZXF1ZXN0cyI6IFsiMSJdLCAiVXNlci1BZ2VudCI6IFsiQ3VzdG9tIFVzZXIgQWdlbnQgU3RyaW5nIl0sICJWaWEiOiBbIjEuMSAwOGYzMjNkZWFkYmVlZmE3YWYzNGQ1ZmViNDE0Y2UyNy5jbG91ZGZyb250Lm5ldCAoQ2xvdWRGcm9udCkiXSwgIlgtQW16LUNmLUlkIjogWyJjRGVoVlFvWm54NDNWWVFiOWoyLW52Q2gtOXozOTZVaGJwMDI3WTJKdmtDUE5MbUdKSHFsYUE9PSJdLCAiWC1Gb3J3YXJkZWQtRm9yIjogWyIxMjcuMC4wLjEsIDEyNy4wLjAuMiJdLCAiWC1Gb3J3YXJkZWQtUG9ydCI6IFsiNDQzIl0sICJYLUZvcndhcmRlZC1Qcm90byI6IFsiaHR0cHMiXX0sICJyZXF1ZXN0Q29udGV4dCI6IHsiYWNjb3VudElkIjogIjEyMzQ1Njc4OTAxMiIsICJyZXNvdXJjZUlkIjogIjEyMzQ1NiIsICJzdGFnZSI6ICJwcm9kIiwgInJlcXVlc3RJZCI6ICJjNmFmOWFjNi03YjYxLTExZTYtOWE0MS05M2U4ZGVhZGJlZWYiLCAicmVxdWVzdFRpbWUiOiAiMDkvQXByLzIwMTU6MTI6MzQ6NTYgKzAwMDAiLCAicmVxdWVzdFRpbWVFcG9jaCI6IDE0Mjg1ODI4OTYwMDAsICJpZGVudGl0eSI6IHsiY29nbml0b0lkZW50aXR5UG9vbElkIjogbnVsbCwgImFjY291bnRJZCI6IG51bGwsICJjb2duaXRvSWRlbnRpdHlJZCI6IG51bGwsICJjYWxsZXIiOiBudWxsLCAiYWNjZXNzS2V5IjogbnVsbCwgInNvdXJjZUlwIjogIjEyNy4wLjAuMSIsICJjb2duaXRvQXV0aGVudGljYXRpb25UeXBlIjogbnVsbCwgImNvZ25pdG9BdXRoZW50aWNhdGlvblByb3ZpZGVyIjogbnVsbCwgInVzZXJBcm4iOiBudWxsLCAidXNlckFnZW50IjogIkN1c3RvbSBVc2VyIEFnZW50IFN0cmluZyIsICJ1c2VyIjogbnVsbH0sICJwYXRoIjogIi9wcm9kL3BhdGgvdG8vcmVzb3VyY2UiLCAicmVzb3VyY2VQYXRoIjogIi97cHJveHkrfSIsICJodHRwTWV0aG9kIjogIlBPU1QiLCAiYXBpSWQiOiAiMTIzNDU2Nzg5MCIsICJwcm90b2NvbCI6ICJIVFRQLzEuMSJ9fSwgInJlc3BvbnNlIjogWzIwMSwgIntcIm1lc3NhZ2VcIjogXCJzdWNjZXNzXCJ9Il19')
212+
55213
56214if __name__ == '__main__': # pragma: no cover
57215 unittest.main() # pragma: no cover
0 commit comments