Skip to content

Commit 1f9fde3

Browse files
rileyfiretailrileyfiretail
authored
Merge pull request #4 from FireTail-io/documentation
Added Sanitization callback
2 parents 3d1537f + 2a4fa37 commit 1f9fde3

File tree

3 files changed

+212
-11
lines changed

3 files changed

+212
-11
lines changed

‎README.md

Lines changed: 38 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -23,9 +23,11 @@ pip install -U firetail-lambda
2323

2424
Implementing Middleware in lambda function
2525
```python
26-
from firetail_lambda import firetail_handler
26+
from firetail_lambda import firetail_handler, firetail_app
2727

28-
@firetail_handler()
28+
app = firetail_app()
29+
30+
@firetail_handler(app)
2931
def lambda_handler(event, context):
3032
return {
3133
"statusCode": 200,
@@ -36,9 +38,11 @@ def lambda_handler(event, context):
3638
```
3739
Multiple Event handlers
3840
```python
39-
from firetail_lambda import firetail_handler
41+
from firetail_lambda import firetail_handler, firetail_app
42+
43+
app = firetail_app()
4044

41-
@firetail_handler()
45+
@firetail_handler(app)
4246
def lambda_handler(event, context):
4347
return {
4448
"statusCode": 200,
@@ -47,7 +51,7 @@ def lambda_handler(event, context):
4751
})
4852
}
4953

50-
@firetail_handler()
54+
@firetail_handler(app)
5155
def lambda_handler_2(event, context):
5256
return {
5357
"statusCode": 200,
@@ -57,3 +61,32 @@ def lambda_handler_2(event, context):
5761
}
5862
```
5963

64+
Custom Sanitization callback
65+
```python
66+
from firetail_lambda import firetail_handler, firetail_app
67+
68+
def sanitize_payloads(event, response):
69+
new_event = copy.copy(event)
70+
remove_headers = ['authorization','Authorization', 'x-api-key']
71+
if 'headers' in event:
72+
for header in remove_headers:
73+
if header in event['headers']:
74+
del new_event['headers'][header]
75+
if 'multiValueHeaders' in event and header in event['multiValueHeaders']:
76+
del new_event['multiValueHeaders'][header]
77+
78+
return new_event, response
79+
80+
app = firetail_app()
81+
app.sanitization_callback = sanitize_payloads
82+
83+
@firetail_handler(app)
84+
def lambda_handler(event, context):
85+
return {
86+
"statusCode": 200,
87+
"body": json.dumps({
88+
"message": "Hello"
89+
})
90+
}
91+
92+
```

‎firetail_lambda/__init__.py

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,16 @@
55
from .version import __version__
66

77

8-
def firetail_handler(enable_sleeper=False):
8+
class firetail_app:
9+
def __init__(self):
10+
self.sanitization_callback=default_sanitization_callback
11+
self.enable_sleeper=False
12+
self.sanitize_headers = ['authorization', 'x-api-key']
13+
14+
def default_sanitization_callback(event, response):
15+
return event, response
16+
17+
def firetail_handler(self):
918
def decorator(func):
1019
def wrapper_func(*args, **kwargs):
1120
start_time = time.time()
@@ -21,11 +30,12 @@ def wrapper_func(*args, **kwargs):
2130
response = func(*args, **kwargs)
2231

2332
# Create our log payload, and print it
33+
event, response = self.sanitization_callback(event, response)
2434
log_payload = base64.b64encode(json.dumps({"event": event,"response": response}).encode("utf-8")).decode("ascii")
2535
print("firetail:loggingapi:%s" % (log_payload))
2636

2737
## Ensure the execution time is >25ms to give the logs API time to propagate our print() to the extension.
28-
if enable_sleeper:
38+
if self.enable_sleeper:
2939
time.sleep(max(time.time() - start_time + 500/1000, 0))
3040

3141
# Return the response from down the chain

‎tests/test_firetail.py

Lines changed: 162 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,148 @@
11

22

33
import contextlib
4+
import copy
45
import json
56
import time
67
import unittest
78
from io import StringIO
89

9-
from firetail_lambda import firetail_handler
10+
from firetail_lambda import firetail_app, firetail_handler
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

12140
class 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

56214
if __name__ == '__main__': # pragma: no cover
57215
unittest.main() # pragma: no cover

0 commit comments

Comments
 (0)