[confcom] Adding new flag for faster hashing and patching tar file usage

src/confcom/HISTORY.rst

@@ -2,6 +2,10 @@
 
 Release History
 ===============
+0.3.4
+++++++
+* adding faster hashing flag to use buffered reader in dmverity-vhd
+
 0.3.3
 ++++++
 * improving testing standards from pytest recommendations

src/confcom/azext_confcom/_help.py

@@ -81,6 +81,10 @@
           type: boolean
           short-summary: 'When enabled, the generated security policy is printed to the command line instead of injected into the input ARM Template'
 
+        - name: --faster-hashing
+          type: boolean
+          short-summary: 'When enabled, the hashing algorithm used to generate the policy is faster but less memory efficient'
+
     examples:
         - name: Input an ARM Template file to inject a base64 encoded Confidential Container Security Policy into the ARM Template
           text: az confcom acipolicygen --template-file "./template.json"

src/confcom/azext_confcom/_params.py

@@ -121,6 +121,12 @@ def load_arguments(self, _):
             required=False,
             help="Print the generated policy in the terminal",
         )
+        c.argument(
+            "faster_hashing",
+            options_list=("--faster-hashing"),
+            required=False,
+            help="Use buffered image reader for dmverity hashing. This will speed up the hashing process but use much more memory.",
+        )
 
     with self.argument_context("confcom katapolicygen") as c:
         c.argument(

src/confcom/azext_confcom/custom.py

@@ -43,6 +43,7 @@ def acipolicygen_confcom(
     print_policy_to_terminal: bool = False,
     disable_stdio: bool = False,
     print_existing_policy: bool = False,
+    faster_hashing: bool = False,
 ):
 
     if sum(map(bool, [input_path, arm_template, image_name])) != 1:
@@ -57,6 +58,8 @@ def acipolicygen_confcom(
         )
     elif save_to_file and arm_template and not (print_policy_to_terminal or outraw or outraw_pretty_print):
         error_out("Must print policy to terminal when saving to file")
+    elif faster_hashing and tar_mapping_location:
+        error_out("Cannot use --faster-hashing with --tar")
 
     if print_existing_policy or outraw or outraw_pretty_print:
         logger.warning(
@@ -124,7 +127,7 @@ def acipolicygen_confcom(
 
     for count, policy in enumerate(container_group_policies):
         policy.populate_policy_content_for_all_images(
-            individual_image=bool(image_name), tar_mapping=tar_mapping
+            individual_image=bool(image_name), tar_mapping=tar_mapping, faster_hashing=faster_hashing
         )
 
         if validate_sidecar:

src/confcom/azext_confcom/data/internal_config.json

@@ -1,5 +1,5 @@
 {
-    "version": "0.3.3",
+    "version": "0.3.4",
     "hcsshim_config": {
         "maxVersion": "1.0.0",
         "minVersion": "0.0.1"

src/confcom/azext_confcom/kata_proxy.py

@@ -41,7 +41,7 @@ def download_binaries():
         needed_assets = ["genpolicy", "genpolicy.exe"]
         # search for genpolicy in the assets from kata-container releases
         for release in r.json():
-            if release.get("tag_name").startswith("genpolicy"):
+            if "genpolicy" in release.get("tag_name"):
                 # these should be newest to oldest
                 for asset in release["assets"]:
                     # download the file if it contains genpolicy

src/confcom/azext_confcom/os_util.py

@@ -95,7 +95,7 @@ def load_tar_mapping_from_file(path: str) -> dict:
     return raw_json
 
 
-def map_image_from_tar(image_name: str, tar: TarFile, tar_location: str):
+def map_image_from_tar_backwards_compatibility(image_name: str, tar: TarFile, tar_location: str):
     tar_dir = os.path.dirname(tar_location)
     # grab all files in the folder and only take the one that's named with hex values and a json extension
     members = tar.getmembers()
@@ -139,3 +139,41 @@ def map_image_from_tar(image_name: str, tar: TarFile, tar_location: str):
     image_info["Architecture"] = image_info_raw.get("architecture")
 
     return image_info
+
+
+def map_image_from_tar(image_name: str, tar: TarFile, tar_location: str):
+    tar_dir = os.path.dirname(tar_location)
+    info_file = None
+    info_file_name = "manifest.json"
+    # if there's more than one image in the tarball, we need to do some more logic
+    if len(info_file_name) > 0:
+        # extract just the manifest file and see if any of the RepoTags match the image_name we're searching for
+        # the manifest.json should have a list of all the image tags
+        # and what json files they map to to get env vars, startup cmd, etc.
+        tar.extract(info_file_name, path=tar_dir)
+        manifest_path = os.path.join(tar_dir, info_file_name)
+        manifest = load_json_from_file(manifest_path)
+        # if we match a RepoTag to the image, stop searching
+        for image in manifest:
+            if image_name in image.get("RepoTags"):
+                info_file = image.get("Config")
+                break
+        # remove the extracted manifest file to clean up
+        os.remove(manifest_path)
+    else:
+        eprint(f"Tarball at {tar_location} contains no images")
+
+    if not info_file:
+        return None
+    tar.extract(info_file, path=tar_dir)
+
+    # get the path of the json file and read it in
+    image_info_file_path = os.path.join(tar_dir, info_file)
+    image_info_raw = load_json_from_file(image_info_file_path)
+    # delete the extracted json file to clean up
+    os.remove(image_info_file_path)
+    image_info = image_info_raw.get("config")
+    # importing the constant from config.py gives a circular dependency error
+    image_info["Architecture"] = image_info_raw.get("architecture")
+
+    return image_info

src/confcom/azext_confcom/rootfs_proxy.py

@@ -88,7 +88,11 @@ def __init__(self):
             os.chmod(self.policy_bin, st.st_mode | stat.S_IXUSR)
 
     def get_policy_image_layers(
-        self, image: str, tag: str, tar_location: str = ""
+        self,
+        image: str,
+        tag: str,
+        tar_location: str = "",
+        faster_hashing=False
     ) -> List[str]:
         image_name = f"{image}:{tag}"
         # populate layer info
@@ -107,6 +111,9 @@ def get_policy_image_layers(
         else:
             arg_list += ["-d"]
 
+        if not tar_location and faster_hashing:
+            arg_list += ["-b"]
+
         # add the image to the end of the parameter list
         arg_list += ["roothash", "-i", f"{image_name}"]
 

src/confcom/azext_confcom/security_policy.py

@@ -352,7 +352,7 @@ def _policy_serialization(self, pretty_print=False) -> str:
 
     # pylint: disable=R0914, R0915
     def populate_policy_content_for_all_images(
-        self, individual_image=False, tar_mapping=None
+        self, individual_image=False, tar_mapping=None, faster_hashing=False,
     ) -> None:
         # suppress warning which will break the progress bar
         warnings.filterwarnings(
@@ -412,7 +412,7 @@ def populate_policy_content_for_all_images(
                         image.set_command(command)
 
                     # merge envs for user container image
-                    envs = image_info.get("Env")
+                    envs = image_info.get("Env") or []
                     env_names = [
                         env_var[
                             config.POLICY_FIELD_CONTAINERS_ELEMENTS_ENVS_RULE
@@ -471,7 +471,7 @@ def populate_policy_content_for_all_images(
                     tar_location = get_tar_location_from_mapping(tar_mapping, image_name)
                 # populate layer info
                 image.set_layers(proxy.get_policy_image_layers(
-                    image.base, image.tag, tar_location=tar_location if tar else ""
+                    image.base, image.tag, tar_location=tar_location if tar else "", faster_hashing=faster_hashing
                 ))
 
                 progress.update()

src/confcom/azext_confcom/template_util.py

@@ -81,9 +81,14 @@ def get_image_info(progress, message_queue, tar_mapping, image):
         if tar_location:
             with tarfile.open(tar_location) as tar:
                 # get all the info out of the tarfile
-                image_info = os_util.map_image_from_tar(
-                    image_name, tar, tar_location
-                )
+                try:
+                    image_info = os_util.map_image_from_tar_backwards_compatibility(
+                        image_name, tar, tar_location
+                    )
+                except IndexError:
+                    image_info = os_util.map_image_from_tar(
+                        image_name, tar, tar_location
+                    )
                 if image_info is not None:
                     tar = True
                     message_queue.append(f"{image_name} read from local tar file")

src/confcom/azext_confcom/tests/latest/test_confcom_arm.py

@@ -1976,8 +1976,8 @@ def setUpClass(cls):
         temp_policies = load_policy_from_arm_template_str(cls.custom_json, "")
         cls.aci_policy = temp_policies[0]
         cls.aci_policy2 = temp_policies[1]
-        cls.aci_policy.populate_policy_content_for_all_images()
-        cls.aci_policy2.populate_policy_content_for_all_images()
+        cls.aci_policy.populate_policy_content_for_all_images(faster_hashing=True)
+        cls.aci_policy2.populate_policy_content_for_all_images(faster_hashing=True)
 
     def test_multiple_policies(self):
         container_start = "containers := "

src/confcom/azext_confcom/tests/latest/test_confcom_scenario.py

@@ -635,7 +635,7 @@ def test_docker_pull(self):
 
             self.assertEqual(
                 image.id,
-                "sha256:e525c930fe751104ff24c356a7bcfad66ce4b92797780eb38dc2ff95d7a66fdc",
+                "sha256:d49a5025be10344cce77d178103a225cb5d7316861e5d8f106e7ff278ae51b62",
             )
 
     def test_infrastructure_svn(self):

src/confcom/setup.py

@@ -18,7 +18,7 @@
 
     logger.warn("Wheel is not available, disabling bdist_wheel hook")
 
-VERSION = "0.3.3"
+VERSION = "0.3.4"
 
 # The full list of classifiers is available at
 # https://pypi.python.org/pypi?%3Aaction=list_classifiers