diff --git a/flask_lambda.py b/flask_lambda.py index 1a57560..5e0d2e6 100644 --- a/flask_lambda.py +++ b/flask_lambda.py @@ -15,6 +15,9 @@ # under the License. import sys +import logging +import base64 +from werkzeug.wrappers import Response try: from urllib import urlencode @@ -35,12 +38,15 @@ __version__ = '0.0.4' +logger = logging.getLogger() +logger.setLevel(logging.INFO) def make_environ(event): environ = {} + headers = event['headers'] or {} - for hdr_name, hdr_value in event['headers'].items(): + for hdr_name, hdr_value in headers.items(): hdr_name = hdr_name.replace('-', '_').upper() if hdr_name in ['CONTENT_TYPE', 'CONTENT_LENGTH']: environ[hdr_name] = hdr_value @@ -49,6 +55,11 @@ def make_environ(event): http_hdr_name = 'HTTP_%s' % hdr_name environ[http_hdr_name] = hdr_value + if 'HTTP_X_FORWARDED_PORT' not in environ: + environ['HTTP_X_FORWARDED_PORT'] = '443' + if 'HTTP_X_FORWARDED_PROTO' not in environ: + environ['HTTP_X_FORWARDED_PROTO'] = 'https' + qs = event['queryStringParameters'] environ['REQUEST_METHOD'] = event['httpMethod'] @@ -61,18 +72,30 @@ def make_environ(event): environ['SERVER_PORT'] = environ['HTTP_X_FORWARDED_PORT'] environ['SERVER_PROTOCOL'] = 'HTTP/1.1' - environ['CONTENT_LENGTH'] = str( - len(event['body']) if event['body'] else '' - ) + if 'isBase64Encoded' in event and event['isBase64Encoded'] is True: + if 'body' in event and event['body'] is not None and 0 < len(event['body']): + tmp_body = base64.b64decode(event['body']) + environ['CONTENT_LENGTH'] = str(len(tmp_body)) + environ['wsgi.input'] = StringIO(tmp_body.decode('utf-8')) + else: + environ['CONTENT_LENGTH'] = '0' + environ['wsgi.input'] = StringIO('') + else: + environ['CONTENT_LENGTH'] = str( + len(event['body']) if event['body'] else '' + ) + environ['wsgi.input'] = StringIO(event['body'] or '') environ['wsgi.url_scheme'] = environ['HTTP_X_FORWARDED_PROTO'] - environ['wsgi.input'] = StringIO(event['body'] or '') environ['wsgi.version'] = (1, 0) environ['wsgi.errors'] = sys.stderr environ['wsgi.multithread'] = False environ['wsgi.run_once'] = True environ['wsgi.multiprocess'] = False + if 'Content-Type' in headers: + environ['CONTENT_TYPE'] = headers['Content-Type'] + BaseRequest(environ) return environ @@ -84,27 +107,48 @@ def __init__(self): self.response_headers = None def start_response(self, status, response_headers, exc_info=None): + _ = exc_info self.status = int(status[:3]) self.response_headers = dict(response_headers) class FlaskLambda(Flask): def __call__(self, event, context): + global logger if 'httpMethod' not in event: # In this "context" `event` is `environ` and # `context` is `start_response`, meaning the request didn't # occur via API Gateway and Lambda + logger.info('Calling non-lambda flask app') return super(FlaskLambda, self).__call__(event, context) - response = LambdaResponse() + logger.info('Calling lambda flask app for following event: ' + str(event)) + + flask_environment = make_environ(event) + body = Response.from_app(self.wsgi_app, flask_environment) + + if body.data: + if body.mimetype.startswith("text/") or body.mimetype.startswith("application/json"): + response_body = body.get_data(as_text=True) + is_base64_encoded = False + else: + response_body = base64.b64encode(body.data).decode('utf-8') + is_base64_encoded = True + else: + response_body = "" + is_base64_encoded = False + + ret = { + 'statusCode': body.status_code, + 'headers': {}, + 'body': response_body + } - body = next(self.wsgi_app( - make_environ(event), - response.start_response - )) + if is_base64_encoded: + ret['isBase64Encoded'] = "true" - return { - 'statusCode': response.status, - 'headers': response.response_headers, - 'body': body - } + for key, value in body.headers: + ret['headers'][key] = value + + logger.info('Response of lambda app: ' + str(ret)) + return ret diff --git a/requirements.txt b/requirements.txt index 3a7e533..6da527e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1 +1,2 @@ Flask>=0.10 +Werkzeug \ No newline at end of file