Fix issue with allow_mismatch: false flag (#3078)

* Fix issue with `allow_mismatch: false` flag

* An exception will be raised when `allow_mismatch` flag is set to `false`
  while downloading images

ansible/playbooks/roles/repository/files/download-requirements/src/command/crane.py

@@ -1,7 +1,4 @@
-from os import chmod
 from pathlib import Path
-from shutil import move
-from tempfile import mkstemp
 from typing import List
 
 from src.command.command import Command
@@ -40,13 +37,6 @@ def pull(self, image_name: str,
             crane_params.append('--format=legacy')
 
         crane_params.append(image_name)
-
-        tmpfile = mkstemp()
-
-        crane_params.append(tmpfile[1])
+        crane_params.append(str(destination))
 
         self.run(crane_params)
-
-        chmod(tmpfile[1], 0o0644)
-
-        move(tmpfile[1], str(destination))

ansible/playbooks/roles/repository/files/download-requirements/src/crypt.py

@@ -1,6 +1,6 @@
 from hashlib import sha1, sha256
 from pathlib import Path
-from typing import Callable
+from typing import Callable, Dict
 
 
 def get_hash(req_path: Path, algorithm: Callable) -> str:
@@ -27,3 +27,9 @@ def get_sha256(req_path: Path) -> str:
 def get_sha1(req_path: Path) -> str:
     """ For larger files sha1 algorithm is significantly faster than sha256 """
     return get_hash(req_path, sha1)
+
+
+SHA_ALGORITHMS: Dict[str, Callable] = {
+    'sha1': get_sha1,
+    'sha256': get_sha256
+}

ansible/playbooks/roles/repository/files/download-requirements/src/downloader.py

@@ -0,0 +1,94 @@
+import logging
+from os import chmod
+from pathlib import Path
+from shutil import move
+from tempfile import mkstemp
+from typing import Callable, Dict
+
+from src.command.toolchain import Toolchain, TOOLCHAINS
+from src.config import Config, OSArch
+from src.crypt import SHA_ALGORITHMS
+from src.error import CriticalError, ChecksumMismatch
+
+
+class Downloader:
+    """
+    Wrapper for downloading requirements with checksum validation.
+    """
+
+    def __init__(self, requirements: Dict[str, Dict],
+                 sha_algorithm: str,
+                 download_func: Callable,
+                 download_func_args: Dict = None):
+        """
+        :param requirements: data from parsed requirements file
+        :param sha_algorithm: which algorithm will be used in validating the requirements
+        :param download_func: back end function used for downloading the requirements
+        :param download_func_args: optional args passed to the `download_func`
+        """
+        self.__requirements: Dict[str, Dict] = requirements
+        self.__sha_algorithm: str = sha_algorithm
+        self.__download: Callable = download_func
+        self.__download_args: Dict = download_func_args or {}
+
+    def __is_checksum_valid(self, requirement: str, requirement_file: Path) -> bool:
+        """
+        Check if checksum matches with `requirement` and downloaded file `requirement_file`.
+
+        :param requirement: an entry from the requirements corresponding to the downloaded file
+        :param requirement_file: existing requirement file
+        :returns: True - checksum ok, False - otherwise
+        :raises:
+            :class:`ChecksumMismatch`: can be raised on failed checksum and missing allow_mismatch flag
+        """
+        req = self.__requirements[requirement]
+        if req[self.__sha_algorithm] != SHA_ALGORITHMS[self.__sha_algorithm](requirement_file):
+            try:
+                if req['allow_mismatch']:
+                    logging.warning(f'{requirement} - allow_mismatch flag used')
+                    return True
+
+                return False
+            except KeyError:
+                return False
+
+        return True
+
+    def download(self, requirement: str, requirement_file: Path, sub_key: str = None):
+        """
+        Download `requirement` as `requirement_file` and compare checksums.
+
+        :param requirement: an entry from the requirements corresponding to the downloaded file
+        :param requirement_file: existing requirement file
+        :param sub_key: optional key for the `requirement` such as `url`
+        :raises:
+            :class:`ChecksumMismatch`: can be raised on failed checksum
+        """
+        req = self.__requirements[requirement][sub_key] if sub_key else requirement
+
+        if requirement_file.exists():
+
+            if self.__is_checksum_valid(requirement, requirement_file):
+                logging.debug(f'- {requirement} - checksum ok, skipped')
+                return
+            else:
+                logging.info(f'- {requirement}')
+
+                tmpfile = Path(mkstemp()[1])
+                chmod(tmpfile, 0o0644)
+
+                self.__download(req, tmpfile, **self.__download_args)
+
+                if not self.__is_checksum_valid(requirement, tmpfile):
+                    tmpfile.unlink()
+                    raise ChecksumMismatch(f'- {requirement}')
+
+                move(str(tmpfile), str(requirement_file))
+                return
+
+        logging.info(f'- {requirement}')
+        self.__download(req, requirement_file, **self.__download_args)
+
+        if not self.__is_checksum_valid(requirement, requirement_file):
+            requirement_file.unlink()
+            raise ChecksumMismatch(f'- {requirement}')

ansible/playbooks/roles/repository/files/download-requirements/src/mode/base_mode.py

@@ -7,8 +7,9 @@
 from poyo import parse_string, PoyoException
 
 from src.command.toolchain import Toolchain, TOOLCHAINS
-from src.config import Config
-from src.crypt import get_sha1, get_sha256
+from src.config import Config, OSArch
+from src.crypt import SHA_ALGORITHMS
+from src.downloader import Downloader
 from src.error import CriticalError, ChecksumMismatch
 
 
@@ -131,61 +132,34 @@ def __download_files(self, files: Dict[str, Dict], dest: Path):
         :param files: to be downloaded
         :param dest: where to save the files
         """
-        for file in files:
-            try:
-                filepath = dest / file.split('/')[-1]
-                if files[file]['sha256'] == get_sha256(filepath):
-                    logging.debug(f'- {file} - checksum ok, skipped')
-                    continue
-
-                logging.info(f'- {file}')
-                self._download_file(file, filepath)
-
-                if files[file]['sha256'] != get_sha256(filepath):
-                    raise ChecksumMismatch(f'- {file}')
-
-            except CriticalError:
-                logging.warning(f'Could not download file: {file}')
+        downloader: Downloader = Downloader(files, 'sha256', self._download_file)
+        for req_file in files:
+            filepath = dest / req_file.split('/')[-1]
+            downloader.download(req_file, filepath)
 
     def __download_grafana_dashboards(self):
         """
         Download grafana dashboards under `self._requirements['grafana-dashboards']`.
         """
         dashboards: Dict[str, Dict] = self._requirements['grafana-dashboards']
+        downloader: Downloader = Downloader(dashboards, 'sha256', self._download_grafana_dashboard)
         for dashboard in dashboards:
-            try:
-                output_file = self._cfg.dest_grafana_dashboards / f'{dashboard}.json'
-
-                if dashboards[dashboard]['sha256'] == get_sha256(output_file):
-                    logging.debug(f'- {dashboard} - checksum ok, skipped')
-                    continue
-
-                logging.info(f'- {dashboard}')
-                self._download_grafana_dashboard(dashboards[dashboard]['url'], output_file)
-
-                if dashboards[dashboard]['sha256'] != get_sha256(output_file):
-                    raise ChecksumMismatch(f'- {dashboard}')
-
-            except CriticalError:
-                logging.warning(f'Could not download grafana dashboard: {dashboard}')
+            output_file = self._cfg.dest_grafana_dashboards / f'{dashboard}.json'
+            downloader.download(dashboard, output_file, 'url')
 
     def __download_crane(self):
         """
         Download Crane package if needed and setup it's environment.
         """
         crane_path = self._cfg.dest_dir / 'crane'
         crane_package_path = Path(f'{crane_path}.tar.gz')
-
         cranes = self._requirements['cranes']
         first_crane = next(iter(cranes))  # right now we use only single crane source
-        if cranes[first_crane]['sha256'] == get_sha256(crane_package_path):
-            logging.debug('crane - checksum ok, skipped')
-        else:
-            self._download_crane_binary(first_crane, crane_package_path)
 
-            if cranes[first_crane]['sha256'] != get_sha256(crane_package_path):
-                raise ChecksumMismatch('crane')
+        downloader: Downloader = Downloader(cranes, 'sha256', self._download_crane_binary)
+        downloader.download(first_crane, crane_package_path)
 
+        if not crane_path.exists():
             self._tools.tar.unpack(crane_package_path, Path('crane'), directory=self._cfg.dest_dir)
             chmod(crane_path, 0o0755)
 
@@ -199,31 +173,18 @@ def _download_images(self):
         """
         Download images under `self._requirements['images']` using Crane.
         """
-        platform: str = 'linux/amd64' if self._cfg.os_arch.X86_64 else 'linux/arm64'
-        images = self._requirements['images']
-        for image in images:
-            try:
-                url, version = image.split(':')
-                filename = Path(f'{url.split("/")[-1]}-{version}.tar')  # format: image_version.tar
-
-                image_file = self._cfg.dest_images / filename
-                if images[image]['sha1'] == get_sha1(image_file):
-                    logging.debug(f'- {image} - checksum ok, skipped')
-                    continue
-
-                logging.info(f'- {image}')
-                self._tools.crane.pull(image, image_file, platform)
-
-                if images[image]['sha1'] != get_sha1(image_file):
-                    try:
-                        if images[image]['allow_mismatch']:
-                            logging.warning(f'- {image} - allow_mismatch flag used, continue downloading')
-                            continue
-                    except KeyError:
-                        raise ChecksumMismatch(f'- {image}')
-
-            except CriticalError:
-                logging.warning(f'Could not download image: `{image}`')
+        platform: str = 'linux/amd64' if self._cfg.os_arch == OSArch.X86_64 else 'linux/arm64'
+        downloader: Downloader = Downloader(self._requirements['images'],
+                                            'sha1',
+                                            self._tools.crane.pull,
+                                            {'platform': platform})
+
+        for image in self._requirements['images']:
+            url, version = image.split(':')
+            filename = Path(f'{url.split("/")[-1]}-{version}.tar')  # format: image_version.tar
+
+            image_file = self._cfg.dest_images / filename
+            downloader.download(image, image_file)
 
     def _cleanup(self):
         """

ansible/playbooks/roles/repository/files/download-requirements/src/mode/debian_family_mode.py

@@ -4,7 +4,7 @@
 import os
 
 from src.config import Config
-from src.mode.base_mode import BaseMode, get_sha256
+from src.mode.base_mode import BaseMode, SHA_ALGORITHMS
 
 
 class DebianFamilyMode(BaseMode):
@@ -85,7 +85,7 @@ def _download_packages(self):
                                    pkg_file.name.startswith(f'{package_info["Package"]}_') and
                                    version in pkg_file.name][0]
 
-                if get_sha256(found_pkg) == package_info['SHA256']:
+                if SHA_ALGORITHMS['sha256'](found_pkg) == package_info['SHA256']:
                     logging.debug(f'- {package} - checksum ok, skipped')
                     continue
 

ansible/playbooks/roles/repository/files/download-requirements/tests/command/test_crane.py

@@ -6,14 +6,10 @@
 
 def test_interface_pull(mocker):
     ''' Check argument construction for crane pull '''
-    mocker.patch('src.command.crane.chmod', return_value=None)
-    mocker.patch('src.command.crane.mkstemp', return_value=[None, '/tmp/tmpfile'])
-    mocker.patch('src.command.crane.move', return_value=None)
-
     with CommandRunMock(mocker, Crane(1).pull, {'image_name': 'image',
                                                 'destination': Path('/some/place'),
                                                 'platform': 'platform',
                                                 'legacy_format': True,
                                                 'insecure': True}) as call_args:
         assert call_args == ['crane', 'pull', '--insecure', '--platform=platform', '--format=legacy',
-                             'image', '/tmp/tmpfile']
+                             'image', '/some/place']