Skip to content

Commit

Permalink
add 3 checks
Browse files Browse the repository at this point in the history
  • Loading branch information
Sandro Braidotti committed Sep 5, 2024
1 parent 71cf9d4 commit e33c929
Show file tree
Hide file tree
Showing 6 changed files with 136 additions and 0 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
from checkov.common.models.consts import ANY_VALUE
from checkov.terraform.checks.resource.base_resource_value_check import BaseResourceValueCheck
from checkov.common.models.enums import CheckCategories

class SagemakerEndpointConfigurationEndpointNameSpecified(BaseResourceValueCheck):
def __init__(self):
name = "Ensure Amazon SageMaker endpoint has a name specified"
id = "CKV_AWS_990"
supported_resources = ['aws_sagemaker_endpoint_configuration']
categories = [CheckCategories.AI_AND_ML]
super().__init__(name=name, id=id, categories=categories, supported_resources=supported_resources)

def get_inspected_key(self):
return "name"

def get_expected_value(self):
return ANY_VALUE

check = SagemakerEndpointConfigurationEndpointNameSpecified()
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
from checkov.common.models.enums import CheckResult, CheckCategories
from checkov.terraform.checks.resource.base_resource_check import BaseResourceCheck

class SagemakerEndpointConfigurationProductionVariantsSpecified(BaseResourceCheck):
def __init__(self):
name = "Ensure Amazon SageMaker endpoint configuration has at least one production variant specified"
id = "CKV_AWS_991"
supported_resources = ['aws_sagemaker_endpoint_configuration']
categories = [CheckCategories.AI_AND_ML]
super().__init__(name=name, id=id, categories=categories, supported_resources=supported_resources)

def scan_resource_conf(self, conf):
production_variants = conf.get("production_variants")
if production_variants:
if isinstance(production_variants, list):
return CheckResult.PASSED if production_variants else CheckResult.FAILED
elif isinstance(production_variants, dict):
return CheckResult.PASSED if 'variant_name' in production_variants and production_variants['variant_name'] else CheckResult.FAILED
return CheckResult.FAILED

check = SagemakerEndpointConfigurationProductionVariantsSpecified()
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
from checkov.common.models.consts import ANY_VALUE
from checkov.terraform.checks.resource.base_resource_value_check import BaseResourceValueCheck
from checkov.common.models.enums import CheckCategories

class SagemakerNotebookLifecycleConfigSpecified(BaseResourceValueCheck):
def __init__(self):
name = "Ensure Amazon SageMaker notebook instances use lifecycle configurations"
id = "CKV_AWS_992"
supported_resources = ['aws_sagemaker_notebook_instance']
categories = [CheckCategories.GENERAL_SECURITY]
super().__init__(name=name, id=id, categories=categories, supported_resources=supported_resources)

def get_inspected_key(self):
return "lifecycle_config_name"

def get_expected_value(self):
return ANY_VALUE

check = SagemakerNotebookLifecycleConfigSpecified()
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
resource "aws_sagemaker_endpoint_configuration" "name_fail" {
kms_key_arn = aws_kms_key.test.arn
production_variants {
variant_name = "variant-1"
model_name = aws_sagemaker_model.m.name
initial_instance_count = 1
instance_type = "ml.t2.medium"
}
tags = {
Name = "foo"
}
}

resource "aws_sagemaker_endpoint_configuration" "name_pass" {
name = "my-endpoint-config"
kms_key_arn = aws_kms_key.test.arn
production_variants {
variant_name = "variant-1"
model_name = aws_sagemaker_model.m.name
initial_instance_count = 1
instance_type = "ml.t2.medium"
}
tags = {
Name = "foo"
}
}


Original file line number Diff line number Diff line change
@@ -0,0 +1,29 @@
resource "aws_sagemaker_endpoint_configuration" "production_variants_pass_1" {
name = "my-endpoint-config"
production_variants {
variant_name = "variant-1"
model_name = aws_sagemaker_model.model.name
initial_instance_count = 1
instance_type = "ml.t2.medium"
}
}

resource "aws_sagemaker_endpoint_configuration" "production_variants_pass_2" {
name = "my-endpoint-config"
production_variants {
variant_name = "variant-1"
model_name = aws_sagemaker_model.model1.name
initial_instance_count = 1
instance_type = "ml.t2.medium"
}
production_variants {
variant_name = "variant-2"
model_name = aws_sagemaker_model.model2.name
initial_instance_count = 2
instance_type = "ml.m5.large"
}
}

resource "aws_sagemaker_endpoint_configuration" "production_variants_fail" {
name = "my-endpoint-config"
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
resource "aws_sagemaker_notebook_instance" "lifecycle_config_fail" {
name = "my-notebook-instance"
role_arn = aws_iam_role.role.arn
instance_type = "ml.t2.medium"

tags = {
Name = "foo"
}
}

resource "aws_sagemaker_notebook_instance" "lifecycle_config_pass" {
name = "my-notebook-instance"
role_arn = aws_iam_role.role.arn
instance_type = "ml.t2.medium"
lifecycle_config_name = "test_lifecycle"

tags = {
Name = "foo"
}
}

0 comments on commit e33c929

Please sign in to comment.