#!/usr/bin/env bash
#######################################################################################
# Copyright ETSI Contributors and Others.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#######################################################################################

# Create a new VM for installing a k8s cluster and its NSG on Azure. SSH key pair ~/.ssh/id_rsa and ~/.ssh/id_rsa.pub must exist.
# To do this it reads the following environment variables:
# - K8S_IMAGE_NAME: name of the new VM
# - RESOURCE_GROUP: name of the resource-group where the VM will be created
# - VNET_NAME: name of the virtual network when creating a new one or referencing an existing one
# - VIM_MGMT_NET: name or ID of the subnet to which the VM will be connected
# - SOURCE_IMAGE_NAME: name of operating system image used (e.g. "Canonical:0001-com-ubuntu-server-jammy:22_04-lts:latest")
# - K8S_FLAVOR_NAME: the VM size to be created (e.g. "Standard_A2_v2")
# - PRIORITY: "Low", "Regular" or Spot"
function create_azure_vm {
    set -eux
    az vm create --resource-group "${RESOURCE_GROUP}" --name "${K8S_IMAGE_NAME}" --image "${SOURCE_IMAGE_NAME}" --size "${K8S_FLAVOR_NAME}" --vnet-name "${VNET_NAME}" --subnet "${VIM_MGMT_NET}" --public-ip-address "" --admin-username ubuntu --priority "${PRIORITY}"
    export K8S_IP=$(az vm show -d -g "${RESOURCE_GROUP}" -n "${K8S_IMAGE_NAME}" --query privateIps | tr -d \")
    # Add a security group rule
    INTERFACE_ID=$(az vm show --resource-group ${RESOURCE_GROUP} --name ${K8S_IMAGE_NAME} --query networkProfile.networkInterfaces[0].id)
    INTERFACE_ID=${INTERFACE_ID:1:-1}
    SECURITY_GROUP_ID=$(az network nic show --id ${INTERFACE_ID} --query networkSecurityGroup.id)
    SECURITY_GROUP_ID=${SECURITY_GROUP_ID:1:-1}
    SECURITY_GROUP_NAME=$(az resource show --ids ${SECURITY_GROUP_ID} --query name)
    SECURITY_GROUP_NAME=${SECURITY_GROUP_NAME:1:-1}
    az network nsg rule create -n microk8s --nsg-name ${SECURITY_GROUP_NAME} --priority 2000 -g ${RESOURCE_GROUP}  --description "Microk8s port" --protocol TCP --destination-port-ranges 16443
}

# Create a new VM for installing a k8s cluster on GCP. SSH key pair ~/.ssh/id_rsa and ~/.ssh/id_rsa.pub must exist.
# To do this it reads the following environment variables:
# - K8S_IMAGE_NAME: name of the new VM
# - GCP_PROJECT: name of project where the VM will be allocated
# - GCP_ZONE: name of the zone (e.g. "europe-west1-b")
# - VIM_MGMT_NET: name or ID of the subnet to which the VM will be connected
# - GCP_MACHINE_TYPE: machine type used for the instance (e.g. "e2-standard-4", run 'gcloud compute machine-types list')
# - GCP_IMAGE_PROJECT: the Google Cloud project against which all image and image family references will be resolved (e.g. "ubuntu-os-cloud")
# - GCP_IMAGE_FAMILY: the image family for the operating system that the boot disk will be initialized with (e.g. "ubuntu-2204-lts")
# - GCP_DISK_SIZE: disk size ib GB
function create_gcp_vm {
    gcloud compute instances create "${K8S_IMAGE_NAME}" --project="${GCP_PROJECT}" --zone="${GCP_ZONE}" --machine-type="${GCP_MACHINE_TYPE}" \
    --network-interface=network-tier=PREMIUM,subnet=${VIM_MGMT_NET} --maintenance-policy=MIGRATE \
    --image-family=${GCP_IMAGE_FAMILY} --image-project=${GCP_IMAGE_PROJECT} --metadata-from-file=ssh-keys=${HOME}/.ssh/id_rsa.pub \
    --create-disk=auto-delete=yes,boot=yes,image-family=${GCP_IMAGE_FAMILY},image-project=${GCP_IMAGE_PROJECT},device-name=${K8S_IMAGE_NAME},mode=rw,size=${GCP_DISK_SIZE} -q
}

# Install via SSH a microk8s cluster on a VM. Writes a kubeconfig.yaml file in $ROBOT_REPORT_FOLDER path.
# To do this it reads the following environment variables:
# - K8S_IP: IP address of the target machine
function install_remote_microk8s {
    set +e
    ssh -T -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null ubuntu@"${K8S_IP}" 'sudo apt-get update -y && sudo apt-get upgrade -y && sudo reboot'
    sleep 90
    ssh -T -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null ubuntu@${K8S_IP} << EOF 2>&1
set -eux
sudo snap install yq
sudo snap install microk8s --classic
sudo usermod -a -G microk8s ubuntu
newgrp microk8s
sudo microk8s.status --wait-ready
sudo microk8s.enable storage dns
set +eux
EOF

    # Enable MetalLB
    ssh -T -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null ubuntu@${K8S_IP} << 'EOF' 2>&1
set -eux
PRIVATE_IP=$(hostname -I | awk '{print $1}')
echo ${PRIVATE_IP}
sudo microk8s.enable metallb:${PRIVATE_IP}-${PRIVATE_IP}
EOF

    # Update the certificate to allow connections from outside as well
    ssh -T -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null ubuntu@${K8S_IP} << EOF 2>&1
set -aux
sudo sed -i "s/\#MOREIPS/IP.3 = ${K8S_IP}/g" /var/snap/microk8s/current/certs/csr.conf.template
cat /var/snap/microk8s/current/certs/csr.conf.template
EOF

    # Save the credentials
    echo ================================================================
    echo K8s cluster credentials:
    echo ================================================================
    echo
    ssh -T -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null ubuntu@${K8S_IP} \
'sudo microk8s.config' | sed  "s/server: .*/server: https:\/\/${K8S_IP}:16443/g" \
| tee ${ROBOT_REPORT_FOLDER}/kubeconfig.yaml
}

# Create an AKS cluster with one node on Azure using a new subnet. Writes a kubeconfig.yaml file in $ROBOT_REPORT_FOLDER path.
# Required the following environment variables:
# - K8S_IMAGE_NAME: name of the AKS cluster to be created
# - RESOURCE_GROUP: name of the resource-group where the VM will be created
# - VNET_NAME: name of the virtual network when creating a new one or referencing an existing one
# - K8S_FLAVOR_NAME: the VM size to be created (e.g. "Standard_D4as_v4")
# IMPORTANT: required a vnet at least /16 because the it's created a new subnet with a random third byte
function create_azure_aks {
    # Gets first subnet prefix and creates a new one with random number in third byte
    set -eux
    SUBNET_PREFIX=$(az network vnet show --resource-group "${RESOURCE_GROUP}" --name "${VNET_NAME}" --query subnets[0].addressPrefix -o tsv)
    IFS=. read BYTE1 BYTE2 BYTE3 BYTE4 <<< "$SUBNET_PREFIX"
    IFS=/ read BYTE4 MASK <<< "$BYTE4"
    if [ "$MASK" -ge "24" ]; then
        BYTE3=$((100 + ($RANDOM % 50) * 2))
        PREFIX="$BYTE1.$BYTE2.${BYTE3}.$BYTE4/$MASK"
        SUBNET_NAME="aks-${BYTE3}"
        BYTE3=$((BYTE3 + 1))
        CIDR="$BYTE1.$BYTE2.${BYTE3}.$BYTE4/$MASK"
        DNS_IP="$BYTE1.$BYTE2.${BYTE3}.10"
        # Verifies that new subnet does not exist previously
        az network vnet subnet show --resource-group "${RESOURCE_GROUP}" --vnet-name "${VNET_NAME}" --name "$SUBNET_NAME" -o table
        if [ "$?" -ne 0 ]; then
            # Creates the subnet in $VNET_NAME network
            az network vnet subnet create --resource-group "${RESOURCE_GROUP}" --vnet-name "${VNET_NAME}" --name "$SUBNET_NAME" --address-prefixes "${PREFIX}"
            if [ "$?" -eq 0 ]; then
                SUBNET_ID=$(az network vnet subnet show --resource-group OSM-CTIO --vnet-name "${VNET_NAME}" --name "${SUBNET_NAME}" --query id -o tsv)

                # Creates k8s cluster
                az aks create -y --resource-group "${RESOURCE_GROUP}" --name "${K8S_IMAGE_NAME}" --node-count 1 --node-vm-size "${K8S_FLAVOR_NAME}" --dns-service-ip "${DNS_IP}" --network-plugin kubenet --service-cidr "${CIDR}" --vnet-subnet-id "${SUBNET_ID}"
                az aks get-credentials --resource-group "${RESOURCE_GROUP}" --name "${K8S_IMAGE_NAME}" --admin -f ${ROBOT_REPORT_FOLDER}/kubeconfig.yaml
            fi
        fi
    fi
}

# Create a GKE cluster with one node on GCP. Writes a kubeconfig.yaml file in $ROBOT_REPORT_FOLDER path.
# Required the following environment variables:
# - K8S_IMAGE_NAME: name of the GKE cluster to be created
# - GCP_PROJECT: name of project where the cluster will be allocated
# - GCP_ZONE: name of the zone (e.g. "europe-west1-b")
# - GCP_MACHINE_TYPE: machine type used for the instance (e.g. "e2-standard-4", run 'gcloud compute machine-types list')
# - GCP_DISK_SIZE: disk size ib GB
function create_gcp_gke {
    gcloud container clusters create "${K8S_IMAGE_NAME}" --project="${GCP_PROJECT}" --zone="${GCP_ZONE}" --num-nodes=1 --machine-type="${GCP_MACHINE_TYPE}" --disk-size "${GCP_DISK_SIZE}" --enable-ip-alias --image-type "COS_CONTAINERD"
}

# Get kubeconfig.yaml from a GKE cluster on GCP. Writes a kubeconfig.yaml file in $ROBOT_REPORT_FOLDER path.
function get_gke_kubeconfig {
    set -eu -o pipefail
    FILE=${ROBOT_REPORT_FOLDER}/kubeconfig.yaml
    SA=osm-sa
    NAMESPACE=osm
    echo "Creating the Kubernetes Service Account with minimal RBAC permissions."
    kubectl apply -f - <<EOF
apiVersion: v1
kind: Namespace
metadata:
  name: ${NAMESPACE}
---
apiVersion: v1
kind: ServiceAccount
metadata:
  name: ${SA}
  namespace: ${NAMESPACE}
---
apiVersion: rbac.authorization.k8s.io/v1
kind: ClusterRole
metadata:
  name: osm-role
rules:
- apiGroups:
  - ""
  resources:
  - users
  - groups
  - serviceaccounts
  verbs:
  - impersonate
- apiGroups:
  - ""
  resources:
  - pods
  verbs:
  - get
- apiGroups:
  - "authorization.k8s.io"
  resources:
  - selfsubjectaccessreviews
  - selfsubjectrulesreviews
  verbs:
  - create
---
apiVersion: rbac.authorization.k8s.io/v1
kind: ClusterRoleBinding
metadata:
  name: osm-crb
roleRef:
  apiGroup: rbac.authorization.k8s.io
  kind: ClusterRole
  name: cluster-admin
subjects:
- kind: ServiceAccount
  name: ${SA}
  namespace: ${NAMESPACE}
EOF

    # Get the service account token and CA cert.
    SA_SECRET_NAME=$(kubectl get -n ${NAMESPACE} sa/${SA} -o "jsonpath={.secrets[0]..name}")
    # Note: service account token is stored base64-encoded in the secret but must
    # be plaintext in kubeconfig.
    SA_TOKEN=$(kubectl get -n ${NAMESPACE} secrets/${SA_SECRET_NAME} -o "jsonpath={.data['token']}" | base64 ${BASE64_DECODE_FLAG})
    CA_CERT=$(kubectl get -n ${NAMESPACE} secrets/${SA_SECRET_NAME} -o "jsonpath={.data['ca\.crt']}")

    # Extract cluster IP from the current context
    CURRENT_CONTEXT=$(kubectl config current-context)
    CURRENT_CLUSTER=$(kubectl config view -o jsonpath="{.contexts[?(@.name == \"${CURRENT_CONTEXT}\"})].context.cluster}")
    CURRENT_CLUSTER_ADDR=$(kubectl config view -o jsonpath="{.clusters[?(@.name == \"${CURRENT_CLUSTER}\"})].cluster.server}")

    echo "Writing $FILE"
    cat > $FILE <<EOF
apiVersion: v1
clusters:
- cluster:
    certificate-authority-data: ${CA_CERT}
    server: ${CURRENT_CLUSTER_ADDR}
  name: ${CURRENT_CLUSTER}
contexts:
- context:
    cluster: ${CURRENT_CLUSTER}
    user: ${CURRENT_CLUSTER}-${SA}
  name: ${CURRENT_CONTEXT}
current-context: ${CURRENT_CONTEXT}
kind: Config
preferences: {}
users:
- name: ${CURRENT_CLUSTER}-${SA}
  user:
    token: ${SA_TOKEN}
EOF

    echo "Done!"
}

# Name of the new VM for K8s (adds a timestamp)
export K8S_IMAGE_NAME=k8stest$(date '+%Y%m%d%H%M')

# Default USE_PAAS_K8S is "FALSE"
if [ -z "${USE_PAAS_K8S}" ]; then
    USE_PAAS_K8S="FALSE"
fi

# Log new environment variables
mkdir -p ${ROBOT_REPORT_FOLDER}
cat <<EOF > ${ROBOT_REPORT_FOLDER}/k8s_environment.rc
export CLOUD_TYPE="${CLOUD_TYPE}"
export USE_PAAS_K8S="${USE_PAAS_K8S}"
EOF

# Branch by USE_PAAS_K8S and CLOUD_TYPE values
if [ "${USE_PAAS_K8S}" == "FALSE" ]; then
    echo "Creating a new IaaS k8s cluster in ${CLOUD_TYPE}"
    if [ "${CLOUD_TYPE}" == "azure" ]; then
        # Create VM on Azure
        create_azure_vm
    elif [ "${CLOUD_TYPE}" == "gcp" ]; then
        # Create VM on GCP
        create_gcp_vm
    else
        echo "CLOUD_TYPE '${CLOUD_TYPE}' not valid for IaaS"
        exit
    fi
    # Add environment variables
    echo "export K8S_IP=\"${K8S_IP}\"" >> ${ROBOT_REPORT_FOLDER}/k8s_environment.rc
    echo "export K8S_IMAGE_NAME=\"${K8S_IMAGE_NAME}\"" >> ${ROBOT_REPORT_FOLDER}/k8s_environment.rc

    # MicroK8s installation
    install_remote_microk8s
else
    echo "Creating a new PaaS k8s cluster in ${CLOUD_TYPE}"
    if [ "${CLOUD_TYPE}" == "azure" ]; then
	create_azure_aks
        echo "export K8S_IMAGE_NAME=\"${K8S_IMAGE_NAME}\"" >> ${ROBOT_REPORT_FOLDER}/k8s_environment.rc
    elif [ "${CLOUD_TYPE}" == "gcp" ]; then
	create_gcp_gke
	get_gke_kubeconfig
        echo "export K8S_IMAGE_NAME=\"${K8S_IMAGE_NAME}\"" >> ${ROBOT_REPORT_FOLDER}/k8s_environment.rc
    else
        echo "CLOUD_TYPE '${CLOUD_TYPE}' not valid for PaaS"
    fi
fi

# Add K8S_CREDENTIALS environment variable
echo "export K8S_CREDENTIALS=${ROBOT_REPORT_FOLDER}/kubeconfig.yaml" >> ${ROBOT_REPORT_FOLDER}/k8s_environment.rc
echo File with new environment was created at ${ROBOT_REPORT_FOLDER}/k8s_environment.rc