__main__.py: Check docker image string

This PR tests that the docker image string provided to tern is of image:tag or image@digest_type:digest format Closes #519 Signed-off-by: PrajwalM2212 <[email protected]>
tern-tools · Feb 11, 2020 · d18c4e1 · d18c4e1
1 parent 0c1a7b0
commit d18c4e1
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 0 deletions.
diff --git a/tern/__main__.py b/tern/__main__.py
@@ -23,6 +23,8 @@
 
 
 # global logger
+from tern.utils.general import check_image_string
+
 logger = logging.getLogger(constants.logger_name)
 logger.setLevel(logging.DEBUG)
 
@@ -86,6 +88,14 @@ def do_main(args):
         if args.dockerfile:
             run.execute_dockerfile(args)
         if args.docker_image:
+            # Check if the image is of image:tag
+            # or image@digest_type:digest format
+            if not check_image_string(args.docker_image):
+                sys.stderr.write('Error running Tern\n'
+                                 'Please provide docker image '
+                                 'string in image:tag or '
+                                 'image@digest_type:digest format\n')
+                sys.exit(1)
             if general.check_tar(args.docker_image):
                 logger.error("%s", errors.incorrect_raw_option)
             else:

diff --git a/tern/utils/general.py b/tern/utils/general.py
@@ -140,3 +140,11 @@ def check_root():
     if os.getuid() == 0:
         return True
     return False
+
+
+def check_image_string(image_str: str):
+    tag_format = r'.+:.+'
+    digest_format = r'.+@.+:.+'
+    if re.match(tag_format, image_str) or re.match(digest_format, image_str):
+        return True
+    return False
diff --git a/tests/test_util_general.py b/tests/test_util_general.py
@@ -0,0 +1,27 @@
+import unittest
+from tern.utils.general import check_image_string
+
+
+class TestUtilGeneral(unittest.TestCase):
+
+    def testImageString(self):
+        correct_strings = [
+            'image@digest_type:digest',
+            'image:tag',
+            'debian:buster',
+            'golang:1.12-alpine',
+            'p12/test@sha256:737aaa0caf3b8f64baa41ebf78c6cd0c43f34fadccc1275a32b8ab5d5b75c344'
+        ]
+
+        incorrect_strings = [
+            'debian',
+            'image',
+            'debian@sha',
+            'test/v1.56'
+        ]
+
+        for image_str in correct_strings:
+            self.assertTrue(check_image_string(image_str))
+
+        for image_str in incorrect_strings:
+            self.assertFalse(check_image_string(image_str))