Add support for stream body

aws_requests_auth/aws_auth.py

@@ -132,16 +132,34 @@ def get_aws_request_headers(self, r, aws_access_key, aws_secret_access_key, aws_
         # Create payload hash (hash of the request body content). For GET
         # requests, the payload is an empty string ('').
         body = r.body if r.body else bytes()
-        try:
-            body = body.encode('utf-8')
-        except (AttributeError, UnicodeDecodeError):
-            # On py2, if unicode characters in present in `body`,
-            # encode() throws UnicodeDecodeError, but we can safely
-            # pass unencoded `body` to execute hexdigest().
-            #
-            # For py3, encode() will execute successfully regardless
-            # of the presence of unicode data
-            body = body
+
+        # check if the body is a stream like object
+        if body and hasattr(r.body, 'read'):
+            try:
+                # read the entire stream until EOF
+                # this should return a byte array like object
+                body = body.read()
+                try:
+                    # check if it's a byte array like object
+                    body.decode()
+                    # we need to set the request body as the new byte array
+                    # as we already consumed the stream
+                    r.body = body
+                except (UnicodeDecodeError, AttributeError):
+                    raise TypeError('read did not return bytes like object')
+            except OSError as e:
+                raise e
+        else:
+            try:
+                body = body.encode('utf-8')
+            except (AttributeError, UnicodeDecodeError):
+                # On py2, if unicode characters in present in `body`,
+                # encode() throws UnicodeDecodeError, but we can safely
+                # pass unencoded `body` to execute hexdigest().
+                #
+                # For py3, encode() will execute successfully regardless
+                # of the presence of unicode data
+                body = body
 
         payload_hash = hashlib.sha256(body).hexdigest()
 

aws_requests_auth/tests/test_aws_auth.py

@@ -2,6 +2,7 @@
 import hashlib
 import mock
 import sys
+import io
 import unittest
 
 from aws_requests_auth.aws_auth import AWSRequestsAuth
@@ -153,6 +154,35 @@ def test_auth_for_post_with_str_body(self):
 
         }, mock_request.headers)
 
+    def test_auth_for_put_with_stream_body(self):
+        auth = AWSRequestsAuth(aws_access_key='YOURKEY',
+                               aws_secret_access_key='YOURSECRET',
+                               aws_host='search-foo.us-east-1.es.amazonaws.com',
+                               aws_region='us-east-1',
+                               aws_service='es')
+        url = 'http://search-foo.us-east-1.es.amazonaws.com:80/'
+        mock_request = mock.Mock()
+        mock_request.url = url
+        mock_request.method = "PUT"
+        mock_request.body = io.BytesIO(b"binary data: \x00\x01")
+        mock_request.headers = {
+            'Content-Type': 'application/octet-stream',
+        }
+
+        frozen_datetime = datetime.datetime(2016, 6, 18, 22, 4, 5)
+        with mock.patch('datetime.datetime') as mock_datetime:
+            mock_datetime.utcnow.return_value = frozen_datetime
+            auth(mock_request)
+        self.assertEqual({
+            'Authorization': 'AWS4-HMAC-SHA256 Credential=YOURKEY/20160618/us-east-1/es/aws4_request, '
+                             'SignedHeaders=host;x-amz-date, '
+                             'Signature=9add35a4e7b0e6a54a7cebe5d2a60837ea6a19f4724fafc51dcd10f701e60f76',
+            'Content-Type': 'application/octet-stream',
+            'x-amz-date': '20160618T220405Z',
+            'x-amz-content-sha256': hashlib.sha256(io.BytesIO(b"binary data: \x00\x01").read()).hexdigest(),
+
+        }, mock_request.headers)
+
     @unittest.skipIf(
         int(sys.version[0]) > 2,
         'python3 produces a different hash that we\'re comparing.',