#!/usr/bin/env bash

# Helper script to generate an IAM Role needed to install role-based authentication to a KFP service account.
#
# Run as:
# $ ./generate_iam_role ${cluster_arn/cluster_name} ${role_name} ${cluster_region} [optional: ${service_namespace} ${service_account}]
#

CLUSTER_ARN="${1}"
ROLE_NAME="${2}"
CLUSTER_REGION="${3:-us-east-1}"
SERVICE_NAMESPACE="${4:-kubeflow}"
SERVICE_ACCOUNT="${5:-pipeline-runner}"
aws_account=$(aws sts get-caller-identity --query Account --output text)
trust_file="trust.json"
assume_role_file="assume-role.json"

cwd=$(dirname $(realpath $0))

# if using an existing cluster, use the cluster arn to get the region and cluster name
# example, cluster_arn=arn:aws:eks:us-east-1:12345678910:cluster/test
cluster_name=$(echo ${CLUSTER_ARN} | cut -d'/' -f2)

# A function to get the OIDC_ID associated with an EKS cluster
function get_oidc_id {
  # TODO: Ideally this should be based on version compatibility instead of command failure
  eksctl utils associate-iam-oidc-provider --cluster ${cluster_name} --region ${CLUSTER_REGION} --approve
  if [[ $? -ge 1 ]]; then
    eksctl utils associate-iam-oidc-provider --name ${cluster_name} --region ${CLUSTER_REGION} --approve
  fi
  
  local oidc=$(aws eks describe-cluster --name ${cluster_name} --region ${CLUSTER_REGION} --query cluster.identity.oidc.issuer --output text)
  oidc_id=$(echo ${oidc} | rev | cut -d'/' -f1 | rev)
}

# A function that generates an IAM role for the given account, cluster, namespace, region
# Parameter:
#    $1: Name of the trust file to generate.
function create_namespaced_iam_role {
  local trust_file_path="${1}"
  # Check if role already exists
  aws iam get-role --role-name ${ROLE_NAME}
  if [[ $? -eq 0 ]]; then
    echo "A role for this cluster and namespace already exists in this account, assuming SageMaker and AssumeRole access and proceeding."
  else
    echo "IAM Role does not exist, creating a new Role for the cluster"
    aws iam create-role --role-name ${ROLE_NAME} --assume-role-policy-document file://${trust_file_path} --output=text --query "Role.Arn"
    aws iam attach-role-policy --role-name ${ROLE_NAME} --policy-arn arn:aws:iam::aws:policy/AmazonSageMakerFullAccess
    aws iam attach-role-policy --role-name ${ROLE_NAME} --policy-arn arn:aws:iam::aws:policy/AWSRoboMaker_FullAccess

    printf '{
      "Version": "2012-10-17",
      "Statement": [
        {
          "Effect": "Allow",
          "Action": "iam:PassRole",
          "Resource": "*"
        }
      ]
    }' > ${assume_role_file}
    aws iam put-role-policy --role-name ${ROLE_NAME} --policy-name AllowPassRole --policy-document file://${assume_role_file}

    printf '{
      "Version": "2012-10-17",
      "Statement": [
        {
          "Effect": "Allow",
          "Action": "sts:AssumeRole",
          "Resource": "arn:aws:iam::'"${aws_account}"':role/*"
        }
      ]
    }' > ${assume_role_file}
    aws iam put-role-policy --role-name ${ROLE_NAME} --policy-name AllowAssumeRole --policy-document file://${assume_role_file}
  fi
}

# Remove the generated trust file
# Parameter:
#    $1: Name of the trust file to delete.
function delete_generated_file {
  rm "${1}" 
}

echo "Get the OIDC ID for the cluster"
get_oidc_id
echo "Delete the trust json file if it already exists"
delete_generated_file "${trust_file}"
echo "Generate a trust json"
"$cwd"/generate_trust_policy ${CLUSTER_REGION} ${aws_account} ${oidc_id} ${SERVICE_NAMESPACE} ${SERVICE_ACCOUNT} > "${trust_file}"
echo "Create the IAM Role using these values"
create_namespaced_iam_role "${trust_file}"
echo "Cleanup for the next run"
delete_generated_file "${trust_file}"

