Merge pull request #28 from mochic/feature/25_upload-checksum-ext-5

add support for tus upload-checksum header

Author: ifedapoolarewaju

tests/test_request.py

@@ -1,4 +1,6 @@
 import mock
+import base64
+import hashlib
 
 from tusclient import client, request
 from tests import mixin
@@ -10,7 +12,6 @@ def setUp(self):
         self.request = request.TusRequest(self.uploader)
 
     def test_perform(self):
-
         with mock.patch.object(self.request, 'handle') as mock_:
             self.request.handle = mock_
             self.request.perform()
@@ -24,6 +25,25 @@ def test_perform(self):
                 'PATCH', '/files/15acd89eabdf5738ffc',
                 f.read(), headers)
 
+    def test_perform_checksum(self):
+        self.uploader.upload_checksum = True
+        self.request_checksum = request.TusRequest(self.uploader)
+        with mock.patch.object(self.request_checksum, 'handle') as mock_:
+            self.request_checksum.handle = mock_
+            self.request_checksum.perform()
+        with open('LICENSE', 'rb') as f:
+            license = f.read()
+            headers = {
+                'upload-offset': 0,
+                'Content-Type': 'application/offset+octet-stream'
+            }
+            headers.update(self.uploader.headers)
+            headers["upload-checksum"] = "sha1 " + \
+                base64.standard_b64encode(hashlib.sha1(license).digest()).decode("ascii")
+            mock_.request.assert_called_with(
+                'PATCH', '/files/15acd89eabdf5738ffc',
+                license, headers)
+
     def test_close(self):
         with mock.patch.object(self.request, 'handle') as mock_:
             self.request.handle = mock_

tests/test_uploader.py

@@ -152,3 +152,10 @@ def test_upload_retry(self, request_mock):
         with pytest.raises(exceptions.TusCommunicationError):
             self.uploader.upload_chunk()
         self.assertEqual(self.uploader._retried, NUM_OF_RETRIES)
+    
+    @mock.patch('tusclient.uploader.TusRequest')
+    def test_upload_checksum(self, request_mock):
+        self.mock_pycurl(request_mock)
+        self.uploader.upload_checksum = True
+        self.uploader.upload()
+        self.assertEqual(self.uploader.offset, self.uploader.file_size)

tusclient/request.py

@@ -1,5 +1,6 @@
 import http.client
 import errno
+import base64
 from future.moves.urllib.parse import urlparse
 
 from tusclient.exceptions import TusUploadFailed
@@ -40,6 +41,9 @@ def __init__(self, uploader):
         }
         self._request_headers.update(uploader.headers)
         self._content_length = uploader.request_length
+        self._upload_checksum = uploader.upload_checksum
+        self._checksum_algorithm = uploader.checksum_algorithm
+        self._checksum_algorithm_name = uploader.checksum_algorithm_name
         self._response = None
 
     @property
@@ -57,7 +61,16 @@ def perform(self):
             host = '{}://{}'.format(self._url.scheme, self._url.netloc)
             path = self._url.geturl().replace(host, '', 1)
 
-            self.handle.request("PATCH", path, self.file.read(self._content_length), self._request_headers)
+            chunk = self.file.read(self._content_length)
+            if self._upload_checksum:
+                self._request_headers["upload-checksum"] = \
+                    " ".join((
+                        self._checksum_algorithm_name,
+                        base64.b64encode(
+                            self._checksum_algorithm(chunk).digest()
+                        ).decode("ascii"),
+                    ))
+            self.handle.request("PATCH", path, chunk, self._request_headers)
             self._response = self.handle.getresponse()
             self.status_code = self._response.status
             self.response_headers = {k.lower(): v for k, v in self._response.getheaders()}
@@ -74,4 +87,4 @@ def close(self):
         """
         close request handle and end request session
         """
-        self.handle.close()
+        self.handle.close()

tusclient/uploader.py

@@ -7,6 +7,7 @@
 from six import iteritems, b, wraps, MAXSIZE
 from six.moves.urllib.parse import urljoin
 import requests
+import hashlib
 
 from tusclient.exceptions import TusUploadFailed, TusCommunicationError
 from tusclient.request import TusRequest
@@ -76,6 +77,9 @@ class Uploader(object):
             would be used. But you can set your own custom fingerprint module by passing it to the constructor.
         - log_func (<function>):
             A logging function to be passed diagnostic messages during file uploads
+        - upload_checksum (bool):
+            Whether or not to supply the Upload-Checksum header along with each
+            chunk. Defaults to False.
 
     :Constructor Args:
         - file_path (str)
@@ -90,13 +94,16 @@ class Uploader(object):
         - url_storage (Optinal [<tusclient.storage.interface.Storage>])
         - fingerprinter (Optional [<tusclient.fingerprint.interface.Fingerprint>])
         - log_func (Optional [<function>])
+        - upload_checksum (Optional[bool])
     """
     DEFAULT_HEADERS = {"Tus-Resumable": "1.0.0"}
     DEFAULT_CHUNK_SIZE = MAXSIZE
+    CHECKSUM_ALGORITHM_PAIR = ("sha1", hashlib.sha1, )
 
     def __init__(self, file_path=None, file_stream=None, url=None, client=None,
                  chunk_size=None, metadata=None, retries=0, retry_delay=30,
-                 store_url=False, url_storage=None, fingerprinter=None, log_func=None):
+                 store_url=False, url_storage=None, fingerprinter=None, 
+                 log_func=None, upload_checksum=False):
         if file_path is None and file_stream is None:
             raise ValueError("Either 'file_path' or 'file_stream' cannot be None.")
 
@@ -122,6 +129,9 @@ def __init__(self, file_path=None, file_stream=None, url=None, client=None,
         self._retried = 0
         self.retry_delay = retry_delay
         self.log_func = log_func
+        self.upload_checksum = upload_checksum
+        self.__checksum_algorithm_name, self.__checksum_algorithm = \
+            self.CHECKSUM_ALGORITHM_PAIR
 
     # it is important to have this as a @property so it gets
     # updated client headers.
@@ -142,6 +152,19 @@ def headers_as_list(self):
         headers = self.headers
         headers_list = ['{}: {}'.format(key, value) for key, value in iteritems(headers)]
         return headers_list
+    
+    @property
+    def checksum_algorithm(self):
+        """The checksum algorithm to be used for the Upload-Checksum extension. 
+        """
+        return self.__checksum_algorithm
+
+    @property
+    def checksum_algorithm_name(self):
+        """The name of the checksum algorithm to be used for the Upload-Checksum 
+        extension. 
+        """
+        return self.__checksum_algorithm_name
 
     @_catch_requests_error
     def get_offset(self):