def test_framework_version_cpu(cpu):
"""
Check that the framework version in the image tag is the same as the one on a running container.
:param cpu: ECR image URI with "cpu" in the name
"""
image = cpu
if "tensorflow-inference" in image:
pytest.skip(msg="TF inference does not have core tensorflow installed")
tested_framework, tag_framework_version = get_framework_and_version_from_tag(
image)
# Module name is torch
if tested_framework == "pytorch":
tested_framework = "torch"
ctx = Context()
container_name = f"framework-version-{image.split('/')[-1].replace('.', '-').replace(':', '-')}"
_start_container(container_name, image, ctx)
output = _run_cmd_on_container(
container_name,
ctx,
f"import {tested_framework}; print({tested_framework}.__version__)",
executable="python")
if is_canary_context():
assert tag_framework_version in output.stdout.strip()
else:
assert tag_framework_version == output.stdout.strip()
def test_framework_version_gpu(gpu, ec2_connection):
"""
Check that the framework version in the image tag is the same as the one on a running container.
:param gpu: ECR image URI with "gpu" in the name
:param ec2_connection: fixture to establish connection with an ec2 instance
"""
image = gpu
if "tensorflow-inference" in image:
pytest.skip(msg="TF inference does not have core tensorflow installed")
tested_framework, tag_framework_version = get_framework_and_version_from_tag(
image)
# Module name is "torch"
if tested_framework == "pytorch":
tested_framework = "torch"
cmd = f'import {tested_framework}; print({tested_framework}.__version__)'
output = ec2.execute_ec2_training_test(ec2_connection,
image,
cmd,
executable="python")
if is_canary_context():
assert tag_framework_version in output.stdout.strip()
else:
assert tag_framework_version == output.stdout.strip()
def test_framework_version_cpu(image):
"""
Check that the framework version in the image tag is the same as the one on a running container.
This function tests CPU, EIA, and Neuron images.
:param image: ECR image URI
"""
if "gpu" in image:
pytest.skip("GPU images will have their framework version tested in test_framework_and_cuda_version_gpu")
image_repo_name, _ = get_repository_and_tag_from_image_uri(image)
if re.fullmatch(r"(pr-|beta-|nightly-)?tensorflow-inference(-eia)?", image_repo_name):
pytest.skip(msg="TF inference for CPU/GPU/EIA does not have core tensorflow installed")
tested_framework, tag_framework_version = get_framework_and_version_from_tag(image)
# Framework name may include huggingface
tested_framework = tested_framework.lstrip("huggingface_")
# Module name is torch
if tested_framework == "pytorch":
tested_framework = "torch"
ctx = Context()
container_name = get_container_name("framework-version", image)
start_container(container_name, image, ctx)
output = run_cmd_on_container(
container_name, ctx, f"import {tested_framework}; print({tested_framework}.__version__)", executable="python"
)
if is_canary_context():
assert tag_framework_version in output.stdout.strip()
else:
assert tag_framework_version == output.stdout.strip()
def test_framework_version_cpu(image):
"""
Check that the framework version in the image tag is the same as the one on a running container.
This function tests CPU, EIA, and Neuron images.
:param image: ECR image URI
"""
if "gpu" in image:
pytest.skip(
"GPU images will have their framework version tested in test_framework_and_cuda_version_gpu"
)
if "tensorflow-inference" in image:
pytest.skip(msg="TF inference does not have core tensorflow installed")
tested_framework, tag_framework_version = get_framework_and_version_from_tag(
image)
# Module name is torch
if tested_framework == "pytorch":
tested_framework = "torch"
ctx = Context()
container_name = get_container_name("framework-version", image)
start_container(container_name, image, ctx)
output = run_cmd_on_container(
container_name,
ctx,
f"import {tested_framework}; print({tested_framework}.__version__)",
executable="python")
if is_canary_context():
assert tag_framework_version in output.stdout.strip()
else:
assert tag_framework_version == output.stdout.strip()
def test_framework_version_cpu(image):
"""
Check that the framework version in the image tag is the same as the one on a running container.
This function tests CPU, EIA, and Neuron images.
:param image: ECR image URI
"""
if "gpu" in image:
pytest.skip(
"GPU images will have their framework version tested in test_framework_and_cuda_version_gpu"
)
image_repo_name, _ = get_repository_and_tag_from_image_uri(image)
if re.fullmatch(
r"(pr-|beta-|nightly-)?tensorflow-inference(-eia|-graviton)?",
image_repo_name):
pytest.skip(
msg=
"TF inference for CPU/GPU/EIA does not have core tensorflow installed"
)
tested_framework, tag_framework_version = get_framework_and_version_from_tag(
image)
# Framework name may include huggingface
if tested_framework.startswith('huggingface_'):
tested_framework = tested_framework[len("huggingface_"):]
# Module name is torch
if tested_framework == "pytorch":
tested_framework = "torch"
elif tested_framework == "autogluon":
tested_framework = "autogluon.core"
ctx = Context()
container_name = get_container_name("framework-version", image)
start_container(container_name, image, ctx)
output = run_cmd_on_container(
container_name,
ctx,
f"import {tested_framework}; print({tested_framework}.__version__)",
executable="python")
if is_canary_context():
assert tag_framework_version in output.stdout.strip()
else:
if tested_framework == "autogluon.core":
assert output.stdout.strip().startswith(tag_framework_version)
elif tested_framework == "torch" and Version(
tag_framework_version) >= Version("1.10.0"):
torch_version_pattern = r"{torch_version}(\+cpu)".format(
torch_version=tag_framework_version)
assert re.fullmatch(
torch_version_pattern, output.stdout.strip()
), (f"torch.__version__ = {output.stdout.strip()} does not match {torch_version_pattern}\n"
f"Please specify framework version as X.Y.Z+cpu")
else:
if "neuron" in image:
assert tag_framework_version in output.stdout.strip()
else:
assert tag_framework_version == output.stdout.strip()
stop_and_remove_container(container_name, ctx)
def test_framework_and_cuda_version_gpu(gpu, ec2_connection):
"""
Check that the framework and cuda version in the image tag is the same as the one on a running container.
:param gpu: ECR image URI with "gpu" in the name
:param ec2_connection: fixture to establish connection with an ec2 instance
"""
image = gpu
tested_framework, tag_framework_version = get_framework_and_version_from_tag(
image)
# Framework Version Check #
# Skip framework version test for tensorflow-inference, since it doesn't have core TF installed
if "tensorflow-inference" not in image:
# Framework name may include huggingface
if tested_framework.startswith('huggingface_'):
tested_framework = tested_framework[len("huggingface_"):]
# Module name is "torch"
if tested_framework == "pytorch":
tested_framework = "torch"
elif tested_framework == "autogluon":
tested_framework = "autogluon.core"
cmd = f"import {tested_framework}; print({tested_framework}.__version__)"
output = ec2.execute_ec2_training_test(ec2_connection,
image,
cmd,
executable="python")
if is_canary_context():
assert tag_framework_version in output.stdout.strip()
else:
if tested_framework == "autogluon.core":
version_to_check = "0.3.1" if tag_framework_version == "0.3.2" else tag_framework_version
assert output.stdout.strip().startswith(version_to_check)
elif tested_framework == "torch" and Version(
tag_framework_version) >= Version("1.10.0"):
torch_version_pattern = r"{torch_version}(\+cu\d+)".format(
torch_version=tag_framework_version)
assert re.fullmatch(
torch_version_pattern, output.stdout.strip()
), (f"torch.__version__ = {output.stdout.strip()} does not match {torch_version_pattern}\n"
f"Please specify framework version as X.Y.Z+cuXXX")
else:
assert tag_framework_version == output.stdout.strip()
# CUDA Version Check #
cuda_version = re.search(r"-cu(\d+)-", image).group(1)
# MXNet inference/HF tensorflow inference and Autogluon containers do not currently have nvcc in /usr/local/cuda/bin, so check symlink
if "mxnet-inference" in image or "autogluon" in image or "huggingface-tensorflow-inference" in image:
cuda_cmd = "readlink /usr/local/cuda"
else:
cuda_cmd = "nvcc --version"
cuda_output = ec2.execute_ec2_training_test(
ec2_connection, image, cuda_cmd, container_name="cuda_version_test")
# Ensure that cuda version in tag is in the container
assert cuda_version in cuda_output.stdout.replace(".", "")
def ec2_connection(request, ec2_instance, ec2_key_name, ec2_instance_type,
region):
"""
Fixture to establish connection with EC2 instance if necessary
:param request: pytest test request
:param ec2_instance: ec2_instance pytest fixture
:param ec2_key_name: unique key name
:param ec2_instance_type: ec2_instance_type pytest fixture
:param region: Region where ec2 instance is launched
:return: Fabric connection object
"""
instance_id, instance_pem_file = ec2_instance
region = P3DN_REGION if ec2_instance_type == "p3dn.24xlarge" else region
ip_address = ec2_utils.get_public_ip(instance_id, region=region)
LOGGER.info(f"Instance ip_address: {ip_address}")
user = ec2_utils.get_instance_user(instance_id, region=region)
LOGGER.info(f"Connecting to {user}@{ip_address}")
conn = Connection(
user=user,
host=ip_address,
connect_kwargs={"key_filename": [instance_pem_file]},
connect_timeout=18000,
)
random.seed(f"{datetime.datetime.now().strftime('%Y%m%d%H%M%S%f')}")
unique_id = random.randint(1, 100000)
artifact_folder = f"{ec2_key_name}-{unique_id}-folder"
s3_test_artifact_location = test_utils.upload_tests_to_s3(artifact_folder)
def delete_s3_artifact_copy():
test_utils.delete_uploaded_tests_from_s3(s3_test_artifact_location)
request.addfinalizer(delete_s3_artifact_copy)
conn.run(
f"aws s3 cp --recursive {test_utils.TEST_TRANSFER_S3_BUCKET}/{artifact_folder} $HOME/container_tests"
)
conn.run(
f"mkdir -p $HOME/container_tests/logs && chmod -R +x $HOME/container_tests/*"
)
# Log into ECR if we are in canary context
if test_utils.is_canary_context():
public_registry = test_utils.PUBLIC_DLC_REGISTRY
test_utils.login_to_ecr_registry(conn, public_registry, region)
return conn
def test_framework_and_cuda_version_gpu(gpu, ec2_connection):
"""
Check that the framework and cuda version in the image tag is the same as the one on a running container.
:param gpu: ECR image URI with "gpu" in the name
:param ec2_connection: fixture to establish connection with an ec2 instance
"""
image = gpu
tested_framework, tag_framework_version = get_framework_and_version_from_tag(
image)
# Framework Version Check #
# Skip framework version test for tensorflow-inference, since it doesn't have core TF installed
if "tensorflow-inference" not in image:
# Module name is "torch"
if tested_framework == "pytorch":
tested_framework = "torch"
if tested_framework == "huggingface_pytorch":
tested_framework = "torch"
if tested_framework == "huggingface_tensorflow":
tested_framework = "tensorflow"
cmd = f"import {tested_framework}; print({tested_framework}.__version__)"
output = ec2.execute_ec2_training_test(ec2_connection,
image,
cmd,
executable="python")
if is_canary_context():
assert tag_framework_version in output.stdout.strip()
else:
assert tag_framework_version == output.stdout.strip()
# CUDA Version Check #
cuda_version = re.search(r"-cu(\d+)-", image).group(1)
# MXNet inference containers do not currently have nvcc in /usr/local/cuda/bin, so check symlink
if "mxnet-inference" in image:
cuda_cmd = "readlink /usr/local/cuda"
else:
cuda_cmd = "nvcc --version"
cuda_output = ec2.execute_ec2_training_test(
ec2_connection, image, cuda_cmd, container_name="cuda_version_test")
# Ensure that cuda version in tag is in the container
assert cuda_version in cuda_output.stdout.replace(".", "")
def _get_safety_ignore_list(image_uri):
"""
Get a list of known safety check issue IDs to ignore, if specified in IGNORE_LISTS.
:param image_uri:
:return: <list> list of safety check IDs to ignore
"""
framework = ("mxnet" if "mxnet" in image_uri else
"pytorch" if "pytorch" in image_uri else "tensorflow")
job_type = "training" if "training" in image_uri else "inference-eia" if "eia" in image_uri else "inference"
python_version = "py2" if "py2" in image_uri else "py3"
# TODO: Remove each if condition on each subsequent release
additional_skips = []
if is_canary_context():
if (framework == "tensorflow" and "2.1" in image_uri) or \
(framework == "pytorch" and "1.4" in image_uri) or \
(framework == "pytorch" and job_type == "training" and "1.5" in image_uri):
additional_skips.append('38414')
return IGNORE_SAFETY_IDS.get(framework, {}).get(job_type, {}).get(
python_version, []) + additional_skips
if not (vulnerability_severity.get("CRITICAL")
or vulnerability_severity.get("HIGH")):
return
raise DependencyCheckFailure(
f"Unrecognized CVEs have been reported : {vulnerability_severity}. "
f"Allowed vulnerabilities are {allowed_vulnerabilities or None}. Please see "
f"{dependency_check_report} for more details.")
@pytest.mark.usefixtures("sagemaker", "huggingface")
@pytest.mark.model("N/A")
@pytest.mark.canary("Run dependency tests regularly on production images")
@pytest.mark.parametrize("ec2_instance_type", ["c5.4xlarge"], indirect=True)
@pytest.mark.skipif(
(is_canary_context() and not is_time_for_canary_safety_scan()),
reason=
"Executing test in canaries pipeline during only a limited period of time.",
)
def test_dependency_check_cpu(cpu, ec2_connection, cpu_only,
x86_compatible_only):
_run_dependency_check_test(cpu, ec2_connection)
@pytest.mark.usefixtures("sagemaker", "huggingface")
@pytest.mark.model("N/A")
@pytest.mark.canary("Run dependency tests regularly on production images")
@pytest.mark.parametrize("ec2_instance_type", ["p3.2xlarge"], indirect=True)
@pytest.mark.skipif(
(is_canary_context() and not is_time_for_canary_safety_scan()),
reason=
def test_framework_version_cpu(image):
"""
Check that the framework version in the image tag is the same as the one on a running container.
This function tests CPU, EIA images.
:param image: ECR image URI
"""
if "gpu" in image:
pytest.skip(
"GPU images will have their framework version tested in test_framework_and_cuda_version_gpu"
)
if "neuron" in image:
pytest.skip(
"Neuron images will have their framework version tested in test_framework_and_neuron_sdk_version"
)
image_repo_name, _ = get_repository_and_tag_from_image_uri(image)
if re.fullmatch(
r"(pr-|beta-|nightly-)?tensorflow-inference(-eia|-graviton)?",
image_repo_name):
pytest.skip(
"Non-gpu tensorflow-inference images will be tested in test_tf_serving_version_cpu."
)
tested_framework, tag_framework_version = get_framework_and_version_from_tag(
image)
# Framework name may include huggingface
if tested_framework.startswith('huggingface_'):
tested_framework = tested_framework[len("huggingface_"):]
# Module name is torch
if tested_framework == "pytorch":
tested_framework = "torch"
elif tested_framework == "autogluon":
tested_framework = "autogluon.core"
ctx = Context()
container_name = get_container_name("framework-version", image)
start_container(container_name, image, ctx)
output = run_cmd_on_container(
container_name,
ctx,
f"import {tested_framework}; print({tested_framework}.__version__)",
executable="python")
if is_canary_context():
assert tag_framework_version in output.stdout.strip()
else:
if tested_framework == "autogluon.core":
version_to_check = "0.3.1" if tag_framework_version == "0.3.2" else tag_framework_version
assert output.stdout.strip().startswith(version_to_check)
# Habana v1.2 binary does not follow the X.Y.Z+cpu naming convention
elif "habana" not in image_repo_name:
if tested_framework == "torch" and Version(
tag_framework_version) >= Version("1.10.0"):
torch_version_pattern = r"{torch_version}(\+cpu)".format(
torch_version=tag_framework_version)
assert re.fullmatch(
torch_version_pattern, output.stdout.strip()
), (f"torch.__version__ = {output.stdout.strip()} does not match {torch_version_pattern}\n"
f"Please specify framework version as X.Y.Z+cpu")
else:
if "neuron" in image:
assert tag_framework_version in output.stdout.strip()
if all(_string in image
for _string in ["pytorch", "habana", "synapseai1.3.0"]):
# Habana Pytorch version looks like 1.10.0a0+gitb488e78 for SynapseAI1.3 PT1.10.1 images
pt_fw_version_pattern = r"(\d+(\.\d+){1,2}(-rc\d)?)((a0\+git\w{7}))"
pt_fw_version_match = re.fullmatch(pt_fw_version_pattern,
output.stdout.strip())
# This is desired for PT1.10.1 images
assert pt_fw_version_match.group(1) == "1.10.0"
else:
assert tag_framework_version == output.stdout.strip()
stop_and_remove_container(container_name, ctx)
def conduct_failure_routine(image, image_allowlist,
ecr_image_vulnerability_list,
upgraded_image_vulnerability_list,
s3_bucket_for_storage):
"""
This method conducts the entire process that is supposed to be followed when ECR test fails. It finds all
the fixable and non fixable vulnerabilities and all the packages that can be upgraded and finally invokes
the Auto-Secure lambda for further processing.
:param image: str, image uri
:param image_allowlist: ScanVulnerabilityList, Vulnerabities that are present in the respective allowlist in the DLC git repo.
:param ecr_image_vulnerability_list: ScanVulnerabilityList, Vulnerabities recently detected WITHOUT running apt-upgrade on the originally released image.
:param upgraded_image_vulnerability_list: ScanVulnerabilityList, Vulnerabilites exisiting in the image WITH apt-upgrade run on it.
:param s3_bucket_for_storage: s3 name of the bucket that would be used for saving all the important data that needs to be stored during failure routine.
:return: dict, a dictionary consisting of the entire summary of the steps run within this method.
"""
s3_filename_for_allowlist = save_scan_vulnerability_list_object_to_s3_in_json_format(
image, upgraded_image_vulnerability_list, "allowlist",
s3_bucket_for_storage)
s3_filename_for_current_image_ecr_scan_list = save_scan_vulnerability_list_object_to_s3_in_json_format(
image, ecr_image_vulnerability_list, "current-ecr-scanlist",
s3_bucket_for_storage)
original_filepath_for_allowlist = get_ecr_scan_allowlist_path(image)
edited_files = [{
"s3_filename": s3_filename_for_allowlist,
"github_filepath": original_filepath_for_allowlist
}]
vulnerabilities_fixable_by_upgrade = get_vulnerabilites_fixable_by_upgrade(
image_allowlist, ecr_image_vulnerability_list,
upgraded_image_vulnerability_list)
newly_found_non_fixable_vulnerabilites = upgraded_image_vulnerability_list - image_allowlist
fixable_list = {}
if vulnerabilities_fixable_by_upgrade:
fixable_list = vulnerabilities_fixable_by_upgrade.vulnerability_list
apt_upgrade_list_filename = f"apt-upgrade-list-{test_utils.get_processor_from_image_uri(image)}.txt"
s3_filename_for_apt_upgrade_list = s3_filename_for_allowlist.replace(
"allowlist.json", apt_upgrade_list_filename)
original_filepath_for_apt_upgrade_list = os.path.join(
os.path.dirname(original_filepath_for_allowlist),
apt_upgrade_list_filename)
new_package_list = fixable_list if isinstance(
fixable_list, list) else list(fixable_list.keys())
create_and_save_package_list_to_s3(
original_filepath_for_apt_upgrade_list,
new_package_list,
s3_filename_for_apt_upgrade_list,
s3_bucket_for_storage,
)
edited_files.append({
"s3_filename":
s3_filename_for_apt_upgrade_list,
"github_filepath":
original_filepath_for_apt_upgrade_list
})
newly_found_non_fixable_list = {}
if newly_found_non_fixable_vulnerabilites:
newly_found_non_fixable_list = newly_found_non_fixable_vulnerabilites.vulnerability_list
message_body = {
"edited_files": edited_files,
"fixable_vulnerabilities": fixable_list,
"non_fixable_vulnerabilities": newly_found_non_fixable_list,
}
## TODO: Remove the is_pr_context before merging ##
if test_utils.is_canary_context(
) and test_utils.is_time_for_invoking_ecr_scan_failure_routine_lambda():
_invoke_lambda(
function_name=test_utils.ECR_SCAN_FAILURE_ROUTINE_LAMBDA,
payload_dict=message_body)
return_dict = copy.deepcopy(message_body)
return_dict["s3_filename_for_allowlist"] = s3_filename_for_allowlist
return_dict[
"s3_filename_for_current_image_ecr_scan_list"] = s3_filename_for_current_image_ecr_scan_list
return return_dict
@dataclass
class SafetyPythonEnvironmentVulnerabilityReport:
"""
One of the DataClasses for parsing Safety Report
"""
report: List[SafetyPackageVulnerabilityReport]
def __post_init__(self):
self.report = [
SafetyPackageVulnerabilityReport(**i) for i in self.report
]
@pytest.mark.model("N/A")
@pytest.mark.skipif(is_canary_context(),
reason="Skipping test because it does not run on canary")
def test_safety_file_exists_and_is_valid(image):
"""
Checks if the image has a safety report at the desired location and fails if any of the
packages in the report have failed the safety check.
:param image: str, image uri
"""
repo_name, image_tag = image.split("/")[-1].split(":")
# Make sure this container name doesn't conflict with the safety check test container name
container_name = f"{repo_name}-{image_tag}-safety-file"
# Add null entrypoint to ensure command exits immediately
run(
f"docker run -id "
f"--name {container_name} "
def test_framework_and_cuda_version_gpu(gpu, ec2_connection):
"""
Check that the framework and cuda version in the image tag is the same as the one on a running container.
:param gpu: ECR image URI with "gpu" in the name
:param ec2_connection: fixture to establish connection with an ec2 instance
"""
image = gpu
tested_framework, tag_framework_version = get_framework_and_version_from_tag(
image)
image_repo_name, _ = get_repository_and_tag_from_image_uri(image)
if re.fullmatch(r"(pr-|beta-|nightly-)?tensorflow-inference",
image_repo_name) and Version(
tag_framework_version) == Version("2.6.3"):
pytest.skip(
"Skipping this test for TF 2.6.3 inference as the v2.6.3 version is already on production"
)
# Framework Version Check #
# For tf inference containers, check TF model server version
if re.fullmatch(
r"(pr-|beta-|nightly-)?tensorflow-inference(-eia|-graviton)?",
image_repo_name):
cmd = f"tensorflow_model_server --version"
output = ec2.execute_ec2_training_test(ec2_connection,
image,
cmd,
executable="bash")
assert re.match(rf"TensorFlow ModelServer: {tag_framework_version}(\D+)?", output.stdout), \
f"Cannot find model server version {tag_framework_version} in {output.stdout}"
else:
# Framework name may include huggingface
if tested_framework.startswith('huggingface_'):
tested_framework = tested_framework[len("huggingface_"):]
# Replace the trcomp string as it is extracted from ECR repo name
tested_framework = tested_framework.replace("_trcomp", "")
# Module name is "torch"
if tested_framework == "pytorch":
tested_framework = "torch"
elif tested_framework == "autogluon":
tested_framework = "autogluon.core"
cmd = f"import {tested_framework}; print({tested_framework}.__version__)"
output = ec2.execute_ec2_training_test(ec2_connection,
image,
cmd,
executable="python")
if is_canary_context():
assert tag_framework_version in output.stdout.strip()
else:
if tested_framework == "autogluon.core":
version_to_check = "0.3.1" if tag_framework_version == "0.3.2" else tag_framework_version
assert output.stdout.strip().startswith(version_to_check)
elif tested_framework == "torch" and Version(
tag_framework_version) >= Version("1.10.0"):
torch_version_pattern = r"{torch_version}(\+cu\d+)".format(
torch_version=tag_framework_version)
assert re.fullmatch(
torch_version_pattern, output.stdout.strip()
), (f"torch.__version__ = {output.stdout.strip()} does not match {torch_version_pattern}\n"
f"Please specify framework version as X.Y.Z+cuXXX")
else:
assert tag_framework_version == output.stdout.strip()
# CUDA Version Check #
cuda_version = re.search(r"-cu(\d+)-", image).group(1)
# MXNet inference/HF tensorflow inference and Autogluon containers do not currently have nvcc in /usr/local/cuda/bin, so check symlink
if "mxnet-inference" in image or "autogluon" in image or "huggingface-tensorflow-inference" in image:
cuda_cmd = "readlink /usr/local/cuda"
else:
cuda_cmd = "nvcc --version"
cuda_output = ec2.execute_ec2_training_test(
ec2_connection, image, cuda_cmd, container_name="cuda_version_test")
# Ensure that cuda version in tag is in the container
assert cuda_version in cuda_output.stdout.replace(".", "")
def test_ecr_scan(image, ecr_client, sts_client, region):
"""
Run ECR Scan Tool on an image being tested, and raise Error if vulnerabilities found
1. Start Scan.
2. For 5 minutes (Run DescribeImages):
(We run this for 5 minutes because the Scan is expected to complete in about 2 minutes, though no
analysis has been performed on exactly how long the Scan takes for a DLC image. Therefore we also
have a 3 minute buffer beyond the expected amount of time taken.)
3.1. If imageScanStatus == COMPLETE: exit loop
3.2. If imageScanStatus == IN_PROGRESS or AttributeNotFound(imageScanStatus): continue loop
3.3. If imageScanStatus == FAILED: raise RuntimeError
4. If DescribeImages.imageScanStatus != COMPLETE: raise TimeOutError
5. assert imageScanFindingsSummary.findingSeverityCounts.HIGH/CRITICAL == 0
:param image: str Image URI for image to be tested
:param ecr_client: boto3 Client for ECR
:param sts_client: boto3 Client for STS
:param region: str Name of region where test is executed
"""
test_account_id = sts_client.get_caller_identity().get("Account")
image_account_id = get_account_id_from_image_uri(image)
image_region = get_region_from_image_uri(image)
image_repo_name, original_image_tag = get_repository_and_tag_from_image_uri(image)
additional_image_tags = get_all_the_tags_of_an_image_from_ecr(ecr_client, image)
if not is_image_available_locally(image):
LOGGER.info(f"Image {image} not available locally!! Pulling the image...")
login_to_ecr_registry(Context(), image_account_id, image_region)
run(f"docker pull {image}")
if not is_image_available_locally(image):
raise RuntimeError("Image shown as not available even after pulling")
for additional_tag in additional_image_tags:
image_uri_with_new_tag = image.replace(original_image_tag, additional_tag)
run(f"docker tag {image} {image_uri_with_new_tag}", hide=True)
if image_account_id != test_account_id:
original_image = image
target_image_repo_name = f"beta-{image_repo_name}"
for additional_tag in additional_image_tags:
image_uri_with_new_tag = original_image.replace(original_image_tag, additional_tag)
new_image_uri = ecr_utils.reupload_image_to_test_ecr(image_uri_with_new_tag, target_image_repo_name, region)
if image_uri_with_new_tag == original_image:
image = new_image_uri
minimum_sev_threshold = get_minimum_sev_threshold_level(image)
LOGGER.info(f"Severity threshold level is {minimum_sev_threshold}")
run_scan(ecr_client, image)
scan_results = ecr_utils.get_ecr_image_scan_results(ecr_client, image, minimum_vulnerability=minimum_sev_threshold)
scan_results = ecr_utils.populate_ecr_scan_with_web_scraper_results(image, scan_results)
ecr_image_vulnerability_list = ScanVulnerabilityList(minimum_severity=CVESeverity[minimum_sev_threshold])
ecr_image_vulnerability_list.construct_allowlist_from_ecr_scan_result(scan_results)
remaining_vulnerabilities = ecr_image_vulnerability_list
if not is_image_covered_by_allowlist_feature(image):
if is_canary_context():
pytest.skip("Skipping the test on the canary.")
common_ecr_scan_allowlist = ScanVulnerabilityList(minimum_severity=CVESeverity[minimum_sev_threshold])
common_ecr_scan_allowlist_path = os.path.join(
os.sep, get_repository_local_path(), "data", "common-ecr-scan-allowlist.json"
)
if os.path.exists(common_ecr_scan_allowlist_path):
common_ecr_scan_allowlist.construct_allowlist_from_file(common_ecr_scan_allowlist_path)
remaining_vulnerabilities = remaining_vulnerabilities - common_ecr_scan_allowlist
if remaining_vulnerabilities:
assert not remaining_vulnerabilities.vulnerability_list, (
f"The following vulnerabilities need to be fixed on {image}:\n"
f"{json.dumps(remaining_vulnerabilities.vulnerability_list, indent=4)}"
)
return
upgraded_image_vulnerability_list, image_scan_allowlist = fetch_other_vulnerability_lists(
image, ecr_client, minimum_sev_threshold
)
s3_bucket_name = ECR_SCAN_HELPER_BUCKET
## In case new vulnerabilities (fixable or non-fixable) are found, then conduct failure routine
newly_found_vulnerabilities = ecr_image_vulnerability_list - image_scan_allowlist
# In case there is no new vulnerability but the allowlist is outdated
vulnerabilities_that_can_be_fixed = image_scan_allowlist - upgraded_image_vulnerability_list
if newly_found_vulnerabilities or vulnerabilities_that_can_be_fixed:
failure_routine_summary = conduct_failure_routine(
image,
image_scan_allowlist,
ecr_image_vulnerability_list,
upgraded_image_vulnerability_list,
s3_bucket_name,
)
(
s3_filename_for_fixable_list,
s3_filename_for_non_fixable_list,
) = process_failure_routine_summary_and_store_data_in_s3(failure_routine_summary, s3_bucket_name)
prepend_message = "Found new vulnerabilities in image." if newly_found_vulnerabilities else "Allowlist is outdated."
display_message = prepend_message + " " + (
f"""Found {len(failure_routine_summary["fixable_vulnerabilities"])} fixable vulnerabilites """
f"""and {len(failure_routine_summary["non_fixable_vulnerabilities"])} non fixable vulnerabilites. """
f"""Refer to files s3://{s3_bucket_name}/{s3_filename_for_fixable_list}, s3://{s3_bucket_name}/{s3_filename_for_non_fixable_list}, """
f"""s3://{s3_bucket_name}/{failure_routine_summary["s3_filename_for_current_image_ecr_scan_list"]} and s3://{s3_bucket_name}/{failure_routine_summary["s3_filename_for_allowlist"]}."""
)
if is_canary_context():
LOGGER.error(display_message)
pytest.skip("Skipping the test failure on the canary.")
else:
raise RuntimeError(display_message)
