Added support for binary input and output.

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

requirements.txt

@@ -1 +1,2 @@
 Flask>=0.10
+Werkzeug