Check hyperwallet client response content type

Author: aseveryn-hw

hyperwallet/tests/test_client.py

@@ -45,7 +45,10 @@ def test_receive_non_json_response(self, session_mock):
 
         session_mock.return_value = mock.MagicMock(
             status_code=404,
-            content=data
+            content=data,
+            headers={
+                "Content-Type": "application/json"
+            }
         )
 
         with self.assertRaises(HyperwalletAPIException) as exc:
@@ -68,7 +71,10 @@ def test_receive_valid_json_error_response(self, session_mock):
 
         session_mock.return_value = mock.MagicMock(
             status_code=400,
-            content=json.dumps(data)
+            content=json.dumps(data),
+            headers={
+                "Content-Type": "application/json"
+            }
         )
 
         with self.assertRaises(HyperwalletAPIException) as exc:
@@ -88,7 +94,10 @@ def test_receive_valid_json_response(self, session_mock):
 
         session_mock.return_value = mock.MagicMock(
             status_code=200,
-            content=json.dumps(data)
+            content=json.dumps(data),
+            headers={
+                "Content-Type": "application/json"
+            }
         )
 
         encoded = json.dumps(data)
@@ -100,6 +109,29 @@ def test_receive_valid_json_response(self, session_mock):
             json.loads(encoded)
         )
 
+    @mock.patch('requests.Session.request')
+    def test_receive_json_error_response_when_content_type_is_not_valid(self, session_mock):
+
+        data = {
+            'key': 'value'
+        }
+
+        session_mock.return_value = mock.MagicMock(
+            status_code=200,
+            content=json.dumps(data),
+            headers={
+                "Content-Type": "wrongContentType"
+            }
+        )
+
+        with self.assertRaises(HyperwalletAPIException) as exc:
+            self.client._makeRequest()
+
+        self.assertEqual(
+            exc.exception.message,
+            'Invalid Content-Type specified in Response Header'
+        )
+
 
 if __name__ == '__main__':
     unittest.main()

hyperwallet/tests/test_encryption.py

@@ -9,8 +9,7 @@
 from jwcrypto.common import json_encode
 from hyperwallet.exceptions import HyperwalletException
 from hyperwallet.utils.encryption import Encryption
-from django.core.validators import URLValidator
-from django.core.exceptions import ValidationError
+from six.moves.urllib.parse import urlparse
 
 
 class EncryptionTest(unittest.TestCase):
@@ -158,19 +157,37 @@ def test_should_throw_exception_when_jws_signature_has_expired(self):
         self.assertEqual(exc.exception.message, 'JWS signature has expired, checked by [exp] JWS header')
 
     def __getJwkKeySet(self, location):
-
+        '''
+        Retrieves JWK key data from given location.
+
+        :param location:
+            Location(can be a URL or path to file) of JWK key data. **REQUIRED**
+        :returns:
+            JWK key set found at given location.
+        '''
         try:
-            URLValidator()(location)
-        except ValidationError:
+            url = urlparse(location)
+            if url.scheme and url.netloc and url.path:
+                return requests.get(location).text
+            raise HyperwalletException('Failed to parse url from string = ' + location)
+        except Exception as e:
             if os.path.isfile(location):
                 with open(location) as f:
                     return f.read()
             else:
                 raise HyperwalletException('Wrong JWK key set location path = ' + location)
 
-        return requests.get(location).text
-
     def __findJwkKeyByAlgorithm(self, jwkKeySet, algorithm):
+        '''
+        Finds JWK key by given algorithm.
+
+        :param jwkKeySet:
+            JSON representation of JWK key set. **REQUIRED**
+        :param algorithm:
+            Algorithm of the JWK key to be found in key set. **REQUIRED**
+        :returns:
+            JWK key with given algorithm.
+        '''
 
         try:
             keySet = json.loads(jwkKeySet)

hyperwallet/utils/apiclient.py

@@ -41,7 +41,7 @@ def __init__(self, username, password, server, encryptionData=None):
         # Hyperwallet SDK.
         self.baseHeaders = {
             'User-Agent': 'Hyperwallet Python SDK v{}'.format(__version__),
-            'Accept': 'application/json',
+            'Accept': 'application/jose+json' if self.encrypted else 'application/json',
             'Content-Type': 'application/jose+json' if self.encrypted else 'application/json'
         }
 
@@ -113,6 +113,8 @@ def _makeRequest(self,
         if response.status_code is 204:
             return {}
 
+        self.__checkResponseHeaderContentType(response)
+
         content = response.content
         if hasattr(content, 'decode'):  # Python 2
             content = content.decode('utf-8')
@@ -193,6 +195,20 @@ def doPut(self, partialUrl, data):
             data=json.dumps(data).encode('utf-8')
         )
 
+    def __checkResponseHeaderContentType(self, response):
+        '''
+        Check response header Content-Type.
+
+        :param response:
+            Response to be checked. **REQUIRED**
+        '''
+
+        if response is None:
+            return
+        contentType = response.headers['Content-Type']
+        if (not self.encrypted and contentType != 'application/json') or (self.encrypted and contentType != 'application/jose+json'):
+            raise HyperwalletAPIException('Invalid Content-Type specified in Response Header')
+
     def __getRequestData(self, data):
         '''
         If encryption is enabled try to encrypt request data, otherwise no action required.