PoC: Let user download targets

tests/test_updater.py

@@ -58,6 +58,8 @@
 import errno
 import sys
 import unittest
+import requests
+import timeit
 
 import tuf
 import tuf.exceptions
@@ -1213,6 +1215,97 @@ def test_6_get_one_valid_targetinfo(self):
 
 
 
+  def test_download_target_custom_download_handler(self):
+    # Create temporary directory (destination directory of downloaded targets)
+    # that will be passed as an argument to 'download_target()'.
+    destination_directory = self.make_temp_directory()
+    target_filepaths = \
+      list(self.repository_updater.metadata['current']['targets']['targets'].keys())
+
+    # Get the target info, which is an argument to 'download_target()'.
+    target_filepath = target_filepaths.pop()
+    targetinfo = self.repository_updater.get_one_valid_targetinfo(target_filepath)
+    download_filepath = os.path.join(destination_directory, target_filepath)
+
+    # Test passing a handler with correct signature but incorrect implementation
+    def test_1(param_1):
+      pass
+
+    self.assertRaises(tuf.exceptions.NoWorkingMirrorError,
+        self.repository_updater.download_target, targetinfo, destination_directory,
+        custom_download_handler=test_1)
+
+    # Checks if the file is permanently stored
+    self.assertFalse(os.path.exists(download_filepath))
+
+    # Test passing a handler with incorrect number of parameters
+    def test_2(param_1, param_2):
+      pass
+
+    with self.assertRaises(tuf.exceptions.NoWorkingMirrorError) as cm:
+      self.repository_updater.download_target(targetinfo, destination_directory,
+          custom_download_handler=test_2)
+
+    for _, mirror_error in six.iteritems(cm.exception.mirror_errors):
+      self.assertTrue(isinstance(mirror_error, TypeError))
+
+    # Checks if the file is permanently stored
+    self.assertFalse(os.path.exists(download_filepath))
+
+    # Test passing a handler throwing an unexpected error
+    def test_3(param_1):
+      raise IOError
+
+    with self.assertRaises(tuf.exceptions.NoWorkingMirrorError) as cm:
+      self.repository_updater.download_target(targetinfo, destination_directory,
+          custom_download_handler=test_3)
+
+    for _, mirror_error in six.iteritems(cm.exception.mirror_errors):
+      self.assertTrue(isinstance(mirror_error, IOError))
+
+    # Checks if the file has been successfully downloaded
+    self.assertFalse(os.path.exists(download_filepath))
+
+    # Test: default case.
+    self.repository_updater.download_target(targetinfo, destination_directory)
+
+    # Checks if the file has been successfully downloaded
+    self.assertTrue(os.path.exists(download_filepath))
+
+    # Define a simple custom download handler
+    def download_handler(url):
+      url = six.moves.urllib.parse.unquote(url).replace('\\', '/')
+      temp_file = tempfile.TemporaryFile()
+
+      try:
+        session = requests.Session()
+        with session.get(url, stream=True,
+            timeout=tuf.settings.SOCKET_TIMEOUT) as response:
+
+          # Check response status
+          response.raise_for_status()
+
+          # Read the raw socket response
+          for chunk in response.iter_content(chunk_size=tuf.settings.CHUNK_SIZE):
+            temp_file.write(chunk)
+
+      except Exception:
+        temp_file.close()
+        raise
+
+      else:
+        return temp_file
+
+    # Test passing a custom download function
+    self.repository_updater.download_target(targetinfo, destination_directory,
+        custom_download_handler=download_handler)
+
+    # Checks if the file has been successfully downloaded
+    self.assertTrue(os.path.exists(download_filepath))
+
+
+
+
   def test_6_download_target(self):
     # Create temporary directory (destination directory of downloaded targets)
     # that will be passed as an argument to 'download_target()'.
@@ -1782,10 +1875,7 @@ def verify_target_file(targets_path):
       self.repository_updater._check_hashes(targets_path, file_hashes)
 
     self.repository_updater._get_file('targets.json', verify_target_file,
-        file_type, file_size, download_safely=True).close()
-
-    self.repository_updater._get_file('targets.json', verify_target_file,
-        file_type, file_size, download_safely=False).close()
+        file_type, file_size).close()
 
   def test_13__targets_of_role(self):
     # Test case where a list of targets is given.  By default, the 'targets'

tuf/client/updater.py

@@ -603,9 +603,10 @@ class Updater(object):
       targets that have changed are returns in a list.  From this list, they
       can request a download by calling 'download_target()'.
 
-    download_target(target, destination_directory):
+    download_target(target, destination_directory, custom_download_handler):
       This method performs the actual download of the specified target.  The
-      file is saved to the 'destination_directory' argument.
+      file is saved to the 'destination_directory' argument.  A custom
+      file download function can be used if provided by the user.
 
     remove_obsolete_targets(destination_directory):
       Any files located in 'destination_directory' that were previously
@@ -1313,7 +1314,7 @@ def _soft_check_file_length(self, file_object, trusted_file_length):
 
 
   def _get_target_file(self, target_filepath, file_length, file_hashes,
-      prefix_filename_with_hash):
+      prefix_filename_with_hash, custom_download_handler=None):
     """
     <Purpose>
       Non-public method that safely (i.e., the file length and hash are
@@ -1339,6 +1340,10 @@ def _get_target_file(self, target_filepath, file_length, file_hashes,
         prefixed with hashes (in this case the server uses other means
         to ensure snapshot consistency).
 
+      custom_download_handler:
+        A user provided function performing the actual target file download.
+        If None, tuf.download.safe_download is used.
+
     <Exceptions>
       tuf.exceptions.NoWorkingMirrorError:
         The target could not be fetched. This is raised only when all known
@@ -1370,7 +1375,7 @@ def verify_target_file(target_file_object):
       target_filepath = os.path.join(dirname, target_digest + '.' + basename)
 
     return self._get_file(target_filepath, verify_target_file,
-        'target', file_length, download_safely=True)
+        'target', file_length, custom_download_handler)
 
 
 
@@ -1655,7 +1660,7 @@ def _get_metadata_file(self, metadata_role, remote_filename,
 
 
   def _get_file(self, filepath, verify_file_function, file_type, file_length,
-      download_safely=True):
+      custom_download_handler=None):
     """
     <Purpose>
       Non-public method that tries downloading, up to a certain length, a
@@ -1682,8 +1687,9 @@ def _get_file(self, filepath, verify_file_function, file_type, file_length,
         The expected length, or upper bound, of the target or metadata file to
         be downloaded.
 
-      download_safely:
-        A boolean switch to toggle safe or unsafe download of the file.
+      custom_download_handler:
+        A user provided function performing the actual target file download.
+        If None, tuf.download.safe_download is used.
 
     <Exceptions>
       tuf.exceptions.NoWorkingMirrorError:
@@ -1708,15 +1714,16 @@ def _get_file(self, filepath, verify_file_function, file_type, file_length,
 
     for file_mirror in file_mirrors:
       try:
-        # TODO: Instead of the more fragile 'download_safely' switch, unroll
-        # the function into two separate ones: one for "safe" download, and the
-        # other one for "unsafe" download? This should induce safer and more
-        # readable code.
-        if download_safely:
-          file_object = tuf.download.safe_download(file_mirror, file_length)
+        if custom_download_handler is not None:
+          # When an external download handler is used, file length verification
+          # is not expected. It is performed by verify_file_function()
+          file_object = custom_download_handler(file_mirror)
 
         else:
-          file_object = tuf.download.unsafe_download(file_mirror, file_length)
+          # Ensure the length of the downloaded file matches 'file_length'
+          # exactly even though it will be redundantly verified one more time
+          # by verify_file_function().
+          file_object = tuf.download.safe_download(file_mirror, file_length)
 
         # Verify 'file_object' according to the callable function.
         # 'file_object' is also verified if decompressed above (i.e., the
@@ -3216,7 +3223,7 @@ def updated_targets(self, targets, destination_directory):
 
 
   def download_target(self, target, destination_directory,
-      prefix_filename_with_hash=True):
+      prefix_filename_with_hash=True, custom_download_handler=None):
     """
     <Purpose>
       Download 'target' and verify it is trusted.
@@ -3241,6 +3248,29 @@ def download_target(self, target, destination_directory,
         to ensure snapshot consistency).
         Default is True.
 
+     custom_download_handler:
+        A user provided function performing the actual target file download.
+        In order to comply with the TUF specification, the function implementation
+        should match the following description:
+
+        def download_handler_func(url)
+          <Purpose>
+            Given the 'url' of the desired file,
+            open a connection to 'url', download it, and return the contents
+            of the file.
+
+          <Arguments>
+            url:
+              A URL string that represents the location of the file.
+
+          <Side Effects>
+            A temprorary file object is created to store the contents of 'url'.
+
+          <Returns>
+            A temporary file object that points to the contents of 'url'.
+
+        If None, tuf.download.safe_download is used.
+
     <Exceptions>
       securesystemslib.exceptions.FormatError:
         If 'target' is not properly formatted.
@@ -3298,6 +3328,6 @@ def download_target(self, target, destination_directory,
     # '_get_target_file()' checks every mirror and returns the first target
     # that passes verification.
     target_file_object = self._get_target_file(target_filepath, trusted_length,
-        trusted_hashes, prefix_filename_with_hash)
+        trusted_hashes, prefix_filename_with_hash, custom_download_handler)
 
     securesystemslib.util.persist_temp_file(target_file_object, destination)