setup-workload-identity.sh

#!/bin/bash

set -e

# Color codes for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# Configuration variables - UPDATE THESE
RESOURCE_GROUP="blob-rbac"
STORAGE_ACCOUNT="blobrbacstg"
CONTAINER_NAME="upload-test"
AKS_CLUSTER_NAME="blob-rbac-aks"
NAMESPACE="blob-rbac"
SERVICE_ACCOUNT_NAME="blob-upload-sa"
MANAGED_IDENTITY_NAME="blob-rbac-storage-workload-identity"  # New identity for workload

echo -e "${GREEN}=== Azure Workload Identity Setup for Blob Upload Test ===${NC}"

# Function to check if command exists
command_exists() {
    command -v "$1" >/dev/null 2>&1
}

# Check prerequisites
echo -e "${YELLOW}Checking prerequisites...${NC}"
if ! command_exists az; then
    echo -e "${RED}Azure CLI is not installed. Please install it first.${NC}"
    exit 1
fi

if ! command_exists kubectl; then
    echo -e "${RED}kubectl is not installed. Please install it first.${NC}"
    exit 1
fi

# Get AKS credentials
echo -e "${YELLOW}Getting AKS credentials...${NC}"
az aks get-credentials --resource-group "$RESOURCE_GROUP" --name "$AKS_CLUSTER_NAME"

# Get OIDC issuer URL
echo -e "${YELLOW}Getting OIDC issuer URL...${NC}"
OIDC_ISSUER_RAW=$(az aks show --resource-group "$RESOURCE_GROUP" --name "$AKS_CLUSTER_NAME" --query "oidcIssuerProfile.issuerUrl" -o json | tr -d '\r')

if [ -z "$OIDC_ISSUER_RAW" ] || [ "$OIDC_ISSUER_RAW" = "null" ]; then
    echo -e "${RED}OIDC issuer not found. Make sure OIDC issuer is enabled on your AKS cluster.${NC}"
    echo "To enable OIDC issuer on your AKS cluster, run:"
    echo "az aks update --resource-group \"$RESOURCE_GROUP\" --name \"$AKS_CLUSTER_NAME\" --enable-oidc-issuer"
    exit 1
fi

# Clean up the JSON output (remove quotes)
OIDC_ISSUER=$(echo "$OIDC_ISSUER_RAW" | tr -d '"')

# Validate OIDC issuer URL format
if [[ ! "$OIDC_ISSUER" =~ ^https:// ]]; then
    echo -e "${RED}Invalid OIDC issuer URL format: $OIDC_ISSUER${NC}"
    exit 1
fi

echo "OIDC Issuer: $OIDC_ISSUER"

# Create user-assigned managed identity for workload
echo -e "${YELLOW}Creating user-assigned managed identity for workload...${NC}"
MANAGED_IDENTITY_EXISTS=$(az identity show --name "$MANAGED_IDENTITY_NAME" --resource-group "$RESOURCE_GROUP" --query "id" -o tsv 2>/dev/null | tr -d '\r' || echo "")

if [ -z "$MANAGED_IDENTITY_EXISTS" ]; then
    echo "Creating new managed identity: $MANAGED_IDENTITY_NAME"
    az identity create --name "$MANAGED_IDENTITY_NAME" --resource-group "$RESOURCE_GROUP"
    echo "Waiting for managed identity to propagate..."
    sleep 30
else
    echo "Managed identity already exists: $MANAGED_IDENTITY_NAME"
fi

# Get the client ID of the workload identity
PRINCIPAL_ID=$(az identity show --name "$MANAGED_IDENTITY_NAME" --resource-group "$RESOURCE_GROUP" --query "principalId" -o tsv | tr -d '\r')

if [ -z "$PRINCIPAL_ID" ]; then
    echo -e "${RED}Failed to get client ID for managed identity.${NC}"
    exit 1
fi

echo "Workload Identity Client ID: $PRINCIPAL_ID"

# Get storage account resource ID
echo -e "${YELLOW}Getting storage account information...${NC}"
STORAGE_ACCOUNT_ID=$(az storage account show \
    --resource-group "$RESOURCE_GROUP" \
    --name "$STORAGE_ACCOUNT" \
    --query id -o tsv | tr -d '\r')

if [ -z "$STORAGE_ACCOUNT_ID" ]; then
    echo -e "${RED}Storage account not found. Creating it...${NC}"
    az storage account create \
        --resource-group "$RESOURCE_GROUP" \
        --name "$STORAGE_ACCOUNT" \
        --sku Standard_LRS \
        --kind StorageV2
    
    STORAGE_ACCOUNT_ID=$(az storage account show \
        --resource-group "$RESOURCE_GROUP" \
        --name "$STORAGE_ACCOUNT" \
        --query id -o tsv | tr -d '\r')
fi

# STORAGE_ACCOUNT_ID="${STORAGE_ACCOUNT_ID}/blobServices/default"
echo -e "Storage Account Blob Service ID: $STORAGE_ACCOUNT_ID"

# Create container if it doesn't exist
echo -e "${YELLOW}Creating storage container...${NC}"
az storage container create \
    --name "$CONTAINER_NAME" \
    --account-name "$STORAGE_ACCOUNT" \
    --auth-mode login

# Get both the client ID and principal ID of the workload identity
CLIENT_ID=$(az identity show --name "$MANAGED_IDENTITY_NAME" --resource-group "$RESOURCE_GROUP" --query "clientId" -o tsv | tr -d '\r')

if [ -z "$CLIENT_ID" ] || [ -z "$PRINCIPAL_ID" ]; then
    echo -e "${RED}Failed to get IDs for managed identity.${NC}"
    exit 1
fi

echo "Workload Identity Client ID: $CLIENT_ID"
echo "Workload Identity Principal ID: $PRINCIPAL_ID"

# Assign Storage Blob Data Contributor role to workload identity
echo -e "${YELLOW}Assigning Storage Blob Data Contributor role to workload identity...${NC}"

# Debug: Show the variable values
echo -e "${YELLOW}Debug Information:${NC}"
echo "PRINCIPAL_ID: $PRINCIPAL_ID"
echo "STORAGE_ACCOUNT_ID: $STORAGE_ACCOUNT_ID"
echo ""

# Execute the command
echo -e "${YELLOW}Executing role assignment...${NC}"
az role assignment create --assignee $PRINCIPAL_ID --role "Storage Blob Data Contributor" --scope $STORAGE_ACCOUNT_ID

if [ $? -ne 0 ]; then
    echo -e "${RED}Failed to assign role to workload identity.${NC}"
    exit 1
fi
# Create federated identity credential
echo -e "${YELLOW}Creating federated identity credential...${NC}"
echo -e "Command being executed:"
echo "az identity federated-credential create --name \"blob-upload-federated-credential\" --identity-name \"$MANAGED_IDENTITY_NAME\" --resource-group \"$RESOURCE_GROUP\" --issuer \"$OIDC_ISSUER\" --subject \"system:serviceaccount:$NAMESPACE:$SERVICE_ACCOUNT_NAME\""

# Check if federated credential already exists
FEDERATED_CRED_EXISTS=$(az identity federated-credential show \
    --name "blob-upload-federated-credential" \
    --identity-name "$MANAGED_IDENTITY_NAME" \
    --resource-group "$RESOURCE_GROUP" \
    --audiences "api://AzureADTokenExchange" \
    --query "name" -o tsv 2>/dev/null | tr -d '\r' || echo "")

if [ -z "$FEDERATED_CRED_EXISTS" ]; then
    az identity federated-credential create \
        --name "blob-upload-federated-credential" \
        --identity-name "$MANAGED_IDENTITY_NAME" \
        --resource-group "$RESOURCE_GROUP" \
        --issuer "$OIDC_ISSUER" \
        --subject "system:serviceaccount:$NAMESPACE:$SERVICE_ACCOUNT_NAME" \
        --audiences "api://AzureADTokenExchange"
else
    echo "Federated identity credential already exists."
fi

# Update Kubernetes manifests
echo -e "${YELLOW}Updating Kubernetes manifests...${NC}"
sed -i "s/YOUR_MANAGED_IDENTITY_CLIENT_ID/$PRINCIPAL_ID/g" k8s/deployment.yaml
sed -i "s/YOUR_STORAGE_ACCOUNT_NAME/$STORAGE_ACCOUNT/g" k8s/deployment.yaml

echo -e "${GREEN}=== Setup Complete! ===${NC}"
echo -e "${GREEN}Configuration Summary:${NC}"
echo "Resource Group: $RESOURCE_GROUP"
echo "Storage Account: $STORAGE_ACCOUNT"
echo "Container: $CONTAINER_NAME"
echo "Workload Managed Identity: $MANAGED_IDENTITY_NAME"
echo "Client ID: $PRINCIPAL_ID"
echo "AKS Cluster: $AKS_CLUSTER_NAME"
echo ""
echo -e "${YELLOW}Next steps:${NC}"
echo "1. Build the Docker image: docker build -t blob-rbac:latest ."
echo "2. Load image into kind/minikube: kind load docker-image blob-rbac:latest"
echo "3. Deploy to Kubernetes: kubectl apply -f k8s/deployment.yaml"
echo "4. Check job status: kubectl get jobs -n blob-rbac"
echo "5. View logs: kubectl logs -n blob-rbac job/blob-upload-test-job"