[test] Unskip ECR Scan and report findings (#892)

Author: saimidu

test/dlc_tests/sanity/test_security_check.py

@@ -1,3 +1,5 @@
+import json
+
 from datetime import datetime
 from time import sleep, time
 
@@ -9,6 +11,7 @@
 
 from test.test_utils import get_framework_and_version_from_tag
 from test.test_utils import ecr as ecr_utils
+from test.test_utils.ecr import CVESeverity
 
 
 @pytest.mark.model("N/A")
@@ -50,11 +53,7 @@ def test_ecr_scan(image, ecr_client):
     :param image: str Image URI for image to be tested
     :param ecr_client: boto3 Client for ECR
     """
-    # TODO: Unskip this test for TF 2.4.1 images
-    framework, version = get_framework_and_version_from_tag(image)
-    if framework == "tensorflow" and Version(version) == Version("2.4.1"):
-        pytest.skip("Skip ECR Scan on TF 2.4.1 DLC images")
-
+    minimum_sev_threshold = "HIGH"
     scan_status = None
     start_time = time()
     ecr_utils.start_ecr_image_scan(ecr_client, image)
@@ -68,6 +67,10 @@ def test_ecr_scan(image, ecr_client):
     if scan_status != "COMPLETE":
         raise TimeoutError(f"ECR Scan is still in {scan_status} state. Exiting.")
     severity_counts = ecr_utils.get_ecr_image_scan_severity_count(ecr_client, image)
-    assert not (
-        severity_counts.get("HIGH", 0) or severity_counts.get("CRITICAL", 0)
-    ), f"Found vulnerabilities in image {image}: {str(severity_counts)}"
+    scan_results = ecr_utils.get_ecr_image_scan_results(ecr_client, image, minimum_vulnerability=minimum_sev_threshold)
+    assert all(
+        count == 0 for sev, count in severity_counts.items() if CVESeverity[sev] >= CVESeverity[minimum_sev_threshold]
+    ), (
+        f"Found vulnerabilities in image {image}: {str(severity_counts)}\n"
+        f"Vulnerabilities: {json.dumps(scan_results, indent=4)}"
+    )

test/test_utils/ecr.py

@@ -1,8 +1,19 @@
 import json
 
+from enum import IntEnum
+
 from test.test_utils import get_repository_and_tag_from_image_uri, LOGGER
 
 
+class CVESeverity(IntEnum):
+    UNDEFINED = 0
+    INFORMATIONAL = 1
+    LOW = 2
+    MEDIUM = 3
+    HIGH = 4
+    CRITICAL = 5
+
+
 class ECRScanFailedError(Exception):
     pass
 
@@ -65,3 +76,21 @@ def get_ecr_image_scan_severity_count(ecr_client, image_uri):
     scan_info = ecr_client.describe_image_scan_findings(repositoryName=repository, imageId={"imageTag": tag})
     severity_counts = scan_info["imageScanFindings"]["findingSeverityCounts"]
     return severity_counts
+
+
+def get_ecr_image_scan_results(ecr_client, image_uri, minimum_vulnerability="HIGH"):
+    """
+    Get list of vulnerabilities from ECR image scan results
+    :param ecr_client:
+    :param image_uri:
+    :param minimum_vulnerability: str representing minimum vulnerability level to report in results
+    :return: list<dict> Scan results
+    """
+    repository, tag = get_repository_and_tag_from_image_uri(image_uri)
+    scan_info = ecr_client.describe_image_scan_findings(repositoryName=repository, imageId={"imageTag": tag})
+    scan_findings = [
+        finding
+        for finding in scan_info["imageScanFindings"]["findings"]
+        if CVESeverity[finding["severity"]] >= CVESeverity[minimum_vulnerability]
+    ]
+    return scan_findings