/**

SAGEMAKER NOTEBOOK - IAM POLICIES

This modules covers policies for the SageMaker Notebook. Some of the Policies include:
* S3
* Secret Manager Access
* Assuming SageMaker Role

**/
// -- Assume SageMaker Role
data "aws_iam_policy_document" "sm_assume_role_policy" {
  statement {
    actions = ["sts:AssumeRole"]
    principals {
      type        = "Service"
      identifiers = ["sagemaker.amazonaws.com"]
    }
  }

  statement {
    actions = ["sts:AssumeRole"]
    principals {
      type        = "Service"
      identifiers = ["lambda.amazonaws.com"]
    }
  }
}

data "aws_caller_identity" "current" {}

// Sagamaker Policies -- Base policy
// BASE NOTEBOOK POLICY
// This follows AWS - https://repost.aws/knowledge-center/sagemaker-minimum-permissions
// This has been slimmed down further to only include permissions we need so far
data "aws_iam_policy_document" "sagemaker_base_policy_document" {
  # checkov:skip=CKV_AWS_111:Ensure IAM policies does not allow write access without constraints
  # checkov:skip=CKV_AWS_109:Sagemaker requires to create network interface without the constraints

  // Sagemaker and ECR
  statement {
    actions = [
      "sagemaker:*"
    ]

    resources = [
      // SageMaker Models
      "arn:aws:sagemaker:*:*:model/*",
      // Training jobs
      "arn:aws:logs:*:*:log-group:/aws/sagemaker/TrainingJobs",
      "arn:aws:sagemaker:*:*:training-job/*",
      // Processing Jobs
      "arn:aws:sagemaker:*:*:processing-job/*",
      // Batch Transform Jobs
      "arn:aws:sagemaker:*:*:transform-job/*",
      // Feature Store: feature groups
      "arn:aws:sagemaker:*:*:feature-group/*",
      // Deployment: Endpoints
      "arn:aws:sagemaker:*:*:endpoint/*",
      "arn:aws:sagemaker:*:*:endpoint-config/*",
      // hyperparameter tuninig (automl)
      "arn:aws:sagemaker:*:*:hyper-parameter-tuning-job/*"
    ]
  }

  // Pass Role
  statement {
    actions = [
      "iam:PassRole",
    ]

    resources = [
      "${aws_iam_role.notebook_iam_role.arn}"
    ]

    condition {
      test     = "StringEquals"
      variable = "iam:PassedToService"
      values = [
        "sagemaker.amazonaws.com",
        "glue.amazonaws.com"
      ]
    }
  }

  // Get Role
  statement {
    actions = [
      "iam:GetRole",
    ]

    resources = [
      "${aws_iam_role.notebook_iam_role.arn}"
    ]
  }

  // EC2 Networking
  statement {
    actions = [
      "ec2:DescribeVpcEndpoints",
      "ec2:DescribeDhcpOptions",
      "ec2:DescribeVpcs",
      "ec2:DescribeSubnets",
      "ec2:DescribeSecurityGroups",
      "ec2:DescribeNetworkInterfaces",
    ]

    resources = [
      "*"
    ]
  }

  statement {
    actions = [
      "ec2:DeleteNetworkInterfacePermission",
      "ec2:DeleteNetworkInterface",
      "ec2:CreateNetworkInterfacePermission",
      "ec2:CreateNetworkInterface"
    ]

    resources = [
      "*"
    ]
  }

  // ECR
  statement {
    actions = [
      "ecr:BatchCheckLayerAvailability",
      "ecr:BatchGetImage",
      "ecr:Describe*",
      "ecr:GetAuthorizationToken",
      "ecr:GetDownloadUrlForLayer"
    ]

    resources = [
      "*"
    ]
  }

  # Cloudwatch
  statement {
    actions = [
      "logs:DescribeLogGroups",
      "logs:DescribeLogStreams",
      "logs:GetLogEvents",
      "logs:GetLogRecord",
      "logs:PutLogEvents",
      "logs:PutMetricFilter",
      "logs:FilterLogEvents",
      "logs:CreateLogStream",
      "logs:CreateLogGroup",
      "logs:DescribeMetricFilters",
      "logs:DescribeMetricFilters"
    ]

    resources = [
      "*"
    ]
  }
}

resource "aws_iam_policy" "sagemaker_base_policy" {
  name        = "Sagemaker-${var.environment}-${var.service_name}-base-policy"
  description = "Base Policy for the sagemaker notebook allowing access to Sagemaker APIs and ECR Pull"
  policy      = data.aws_iam_policy_document.sagemaker_base_policy_document.json
  tags        = local.combined_resource_tags
}

// HYPERPARAMETER OPTIMISATION POLICY
// -- Create HyperParameterOptimisation Job Policy
data "aws_iam_policy_document" "hyperparameter_tuning_policy_document" {

  // cloud watch
  statement {
    actions = [
      "cloudwatch:PutMetricData"
    ]
    resources = [
      "*"
    ]
  }

  // Cloud Watch Logging - Restricted to training jobs logs groups
  statement {
    actions = [
      "logs:CreateLogStream",
      "logs:PutLogEvents",
      "logs:CreateLogGroup",
      "logs:DescribeLogStreams"
    ]
    resources = [
      "arn:aws:logs:*:*:log-group:/aws/sagemaker/TrainingJobs"
    ]
  }
}


resource "aws_iam_policy" "hyperparameter_tuning_policy" {
  name        = "Sagemaker-${var.environment}-${var.service_name}-automl-policy"
  description = "AutoML Policy for the sagemaker notebook"
  policy      = data.aws_iam_policy_document.hyperparameter_tuning_policy_document.json
  tags        = local.combined_resource_tags
}


// S3 ACCESS POLICY
// -- policy document
data "aws_iam_policy_document" "s3_access_policy_document" {
  # checkov:skip=CKV_AWS_111:Ensure IAM policies does not allow write access without constraints
  statement {
    actions = [
      "s3:GetObject",
      "s3:PutObject",
      "s3:PutObjectAcl",
      "s3:PutObjectTagging",
      "s3:GetObjectTagging",
      "s3:DeleteObjectTagging"
    ]
    //TODO: should we limit this to just dev-cucumbers?
    resources = [
      "${data.aws_s3_bucket.selected.arn}/*"
    ]
  }

  statement {
    actions = [
      "s3:DeleteObject",
      "s3:DeleteObjectVersion"
    ]
    resources = [
      "${data.aws_s3_bucket.selected.arn}/sagemaker-*"
    ]
  }

  // Feature Store Permissions on S3 (R/W)
  statement {
    actions = [
      "sagemaker:CreateFeatureGroup",
      "sagemaker:ListFeatureGroups",
      "sagemaker:UpdateFeatureGroup",
      "sagemaker:BatchGetRecord",
      "sagemaker:BatchPutMetrics",
      "sagemaker:GetRecord",
      "sagemaker:PutRecord",
      "sagemaker:DeleteRecord",
      "sagemaker:AddTags",
      "sagemaker:TagResource",
    ]
    resources = [
      "*"
    ]
  }

  // Read access
  statement {
    actions = [
      "s3:GetBucketLocation",
      "s3:ListBucket",
      "s3:GetBucketCors",
      "s3:PutBucketCors",
      "s3:GetBucketAcl"
    ]

    resources = [
      data.aws_s3_bucket.selected.arn
    ]

  }
}

// -- iam policy
resource "aws_iam_policy" "s3_acess_policy" {
  name        = "S3-${var.environment}-${var.service_name}-rw-policy"
  description = "RW access to the bucket"
  policy      = data.aws_iam_policy_document.s3_access_policy_document.json
  tags        = local.combined_resource_tags
}

// Secret Manager Access Policies
// -- policy document
data "aws_iam_policy_document" "secrets_manager_policy_document" {
  statement {
    actions = [
      "secretsmanager:GetResourcePolicy",
      "secretsmanager:GetSecretValue",
      "secretsmanager:DescribeSecret",
      "secretsmanager:ListSecretVersionIds"
    ]
    resources = [
      var.code_repository_secret_arn,
      var.snowflake_secret_arn
    ]
  }

  statement {
    actions = [
      "secretsmanager:GetRandomPassword"
    ]

    resources = [
      "*"
    ]
  }
}

// IAM Policy
// -- iam policy
resource "aws_iam_policy" "secret_manager_policy" {
  name        = "SecretsManager-${var.environment}-${var.service_name}-acess-policy"
  description = "SageMaker Notebook Secret Manager Policy"
  policy      = data.aws_iam_policy_document.secrets_manager_policy_document.json
  tags        = local.combined_resource_tags
}


// Feature Store
// -- policy document
data "aws_iam_policy_document" "featurestore_policy_document" {
  statement {
    actions = [
      "s3:PutObject",
      "s3:GetBucketAcl",
      "s3:PutObjectAcl",
      "s3:PutObjectTagging",
      "s3:GetObjectTagging",
      "s3:DeleteObjectTagging"
    ]
    resources = [
      "arn:aws:s3:::*SageMaker*",
      "arn:aws:s3:::*Sagemaker*",
      "arn:aws:s3:::*sagemaker*"
    ]
  }

  statement {
    actions = [
      "s3:GetObject"
    ]

    resources = [
      "arn:aws:s3:::*SageMaker*/metadata/*",
      "arn:aws:s3:::*Sagemaker*/metadata/*",
      "arn:aws:s3:::*sagemaker*/metadata/*"
    ]
  }


  statement {
    actions = [
      "glue:GetTable",
      "glue:UpdateTable"
    ]

    resources = [
      "arn:aws:glue:*:*:catalog",
      "arn:aws:glue:*:*:database/sagemaker_featurestore",
      "arn:aws:glue:*:*:table/sagemaker_featurestore/*"
    ]
  }
}

// IAM Policy
// -- iam policy
resource "aws_iam_policy" "featurestore_policy" {
  name   = "SecretsManager-${var.environment}-${var.service_name}-featurestore-policy"
  policy = data.aws_iam_policy_document.featurestore_policy_document.json
  tags   = local.combined_resource_tags
}

// Model deployment
// -- policy document
data "aws_iam_policy_document" "model_deployment_policy_document" {
  #checkov:skip=CKV_AWS_111:We dont know beforehand the required resources names
  statement {
    actions = [
      "sagemaker:ListEndpoints",
      "sagemaker:ListEndpointConfigs",
      "sagemaker:UpdateEndpoint",
      "sagemaker:UpdateEndpointWeightsAndCapacities",
      "sagemaker:CreateEndpoint",
      "sagemaker:CreateEndpointConfig",
      "sagemaker:DeleteEndpoint",
      "sagemaker:DeleteEndpointConfig",
      "sagemaker:DescribeEndpoint",
      "sagemaker:DescribeEndpointConfig",
      "sagemaker:InvokeEndpoint",
      "sagemaker:CreateModel",
      "sagemaker:DeleteModel",
      "sagemaker:DescribeModel",
      "sagemaker:ListModels",
      "sagemaker:AddTags",
      "sagemaker:TagResource",
    ]
    resources = [
      "*"
    ]
  }

}

// -- iam policy
resource "aws_iam_policy" "model_deployment_policy" {
  name   = "Sagemaker-${var.environment}-${var.service_name}-model-deployment-policy"
  policy = data.aws_iam_policy_document.model_deployment_policy_document.json
  tags   = local.combined_resource_tags
}

data "aws_iam_policy_document" "sagemaker_pipeline_policy_document" {
  #checkov:skip=CKV_AWS_111:We dont know beforehand the required resources names
  statement {
    actions = [
      "sagemaker:CreatePipeline",
      "sagemaker:DescribePipeline",
      "sagemaker:DeletePipeline",
      "sagemaker:ListPipelineExecutionSteps",
      "sagemaker:ListPipelineExecutions",
      "sagemaker:ListPipelineParametersForExecution",
      "sagemaker:ListPipelines",
      "sagemaker:DescribePipelineDefinitionForExecution",
      "sagemaker:DescribePipelineExecution",
      "sagemaker:StartPipelineExecution",
      "sagemaker:StopPipelineExecution",
      "sagemaker:UpdatePipelineExecution",
      "sagemaker:UpdatePipeline",
      "sagemaker:RetryPipelineExecution",
      "sagemaker:SendPipelineExecutionStepFailure",
      "sagemaker:SendPipelineExecutionStepSuccess",
      "sagemaker:CreateTrainingJob",
      "sagemaker:DescribeTrainingJob",
      "sagemaker:UpdateTrainingJob",
      "sagemaker:StopTrainingJob",
      "sagemaker:ListTrainingJobs",
      "sagemaker:CreateProcessingJob",
      "sagemaker:DescribeProcessingJob",
      "sagemaker:CreateModel",
      "sagemaker:CreateModelPackageGroup",
      "sagemaker:CreateModelPackage",
      "sagemaker:ListModelPackages",
      "sagemaker:DeleteModelPackages",
      "sagemaker:DescribeModelPackage",
      "sagemaker:CreateTransformJob",
      "sagemaker:StopTransformJob",
      "sagemaker:DescribeTransformJob",
      "sagemaker:ListTransformJobs",
      "sagemaker:ListTrainingJobsForHyperParameterTuningJob",
      "sagemaker:AddTags",
      "sagemaker:ListAssociations",
      "sagemaker:DeleteModelPackage"
    ]
    resources = [
      "*"
    ]
  }

  statement {
    actions = [
      "lambda:CreateFunction",
      "lambda:UpdateFunctionCode",
      "lambda:InvokeFunction",
      "lambda:DeleteFunction"
    ]

    resources = [
      "arn:aws:lambda:*:*:function:*sagemaker*"
    ]
  }

  statement {
    actions = [
      "iam:PassRole"
    ]

    resources = [
      "arn:aws:iam::*:role/*"
    ]
    condition {
      test     = "StringEquals"
      variable = "iam:PassedToService"
      values = [
        "lambda.amazonaws.com"
      ]
    }
  }

}

resource "aws_iam_policy" "sagemaker_pipeline_policy" {
  name   = "Sagemaker-${var.environment}-${var.service_name}-pipeline-policy"
  policy = data.aws_iam_policy_document.sagemaker_pipeline_policy_document.json
  tags   = local.combined_resource_tags
}