From 3ae5f5d58efcc87a52107f9a229073915420bbe1 Mon Sep 17 00:00:00 2001
From: abrown38 <a.brown@epcc.ed.ac.uk>
Date: Fri, 21 Jun 2024 12:17:11 +0100
Subject: [PATCH] update script for load balancing

---
 register_vms.sh | 84 ++++++++++++++++++++++++++++++++++---------------
 1 file changed, 59 insertions(+), 25 deletions(-)

diff --git a/register_vms.sh b/register_vms.sh
index 69c83ea..9853aa5 100644
--- a/register_vms.sh
+++ b/register_vms.sh
@@ -1,14 +1,16 @@
 #!/bin/bash
 
 # Varialbes Passed By Rundeck
-OS_PROJECT=$1
-WIN_PROJECT=$2
-API_TOKEN=$3
+declare -gr OS_PROJECT=$1
+declare -gr WIN_PROJECT=$2
+declare -gr API_TOKEN=$3
+declare -g RD_JOB_ID=$4
+declare -g RD_EXEC_ID=$5
 
 export OS_CLIENT_CONFIG_FILE=/deployments/$OS_PROJECT-$WIN_PROJECT/clouds.yaml
 
 # check directory exists
-if [ ! -d /deployments/$OS_PROJECT-$WIN_PROJECT/ ]; then
+if [[ ! -d /deployments/$OS_PROJECT-$WIN_PROJECT/ ]]; then
     >&2 echo "Directory Doesn't exist"
     exit 1
 fi
@@ -16,41 +18,73 @@ fi
 RD_JOB_ID=$(cat /deployments/$OS_PROJECT-$WIN_PROJECT/JOB_ID)
 RD_EXEC_ID=$(cat /deployments/$OS_PROJECT-$WIN_PROJECT/EXEC_ID)
 
+initialize_json() {
+    declare -r OUTPUT_FILE="/deployments/$OS_PROJECT-$WIN_PROJECT/data.json"
+
+    jq -n '
+    {
+        "job_id": "'$RD_JOB_ID'",
+        "execution_id": "'$RD_EXEC_ID'",
+        "status": "ok",
+        "DC": {},
+        "WB": {"load_balancer": true, VM: []},
+        "APP": {"load_balancer": false, VM: []},
+        "GPU": {"load_balancer": false, VM: []}
+    }' > $OUTPUT_FILE
+    declare -g -r VM_JSON="/deployments/$OS_PROJECT-$WIN_PROJECT/data.json"
+}
+
+add_vm() {
+
+    declare -r NAME=$1
+    declare -r UUID=$2
+    declare -r IP=$3
+    declare -r TYPE=$4
+
+    if [[ $TYPE == "DC" ]]; then
+        jq ".${TYPE} += {\"name\": \"${NAME}\", \"ip\": \"${IP}\", \"uuid\": \"${UUID}\"}" $VM_JSON > temp.json
+        mv -f temp.json $VM_JSON
+    else
+        jq ".${TYPE}.VM += [{\"name\": \"${NAME}\", \"ip\": \"${IP}\", \"uuid\": \"${UUID}\"}]" $VM_JSON > temp.json
+        mv -f temp.json $VM_JSON
+    fi
+}
 
 post_request() {
-    local INPUT=$1
-    local TYPE=$2
-
-    local UUID=$(echo "$INPUT" | awk '{print $1}')
-    local NAME=$(echo "$INPUT" | awk '{print $2}')
-    local IPv4=$(echo "$INPUT" | awk '{print $4}' | cut -d '=' -f 2)
-
+    #declare -r API_ENDPOINT='https://projects.eidf.ac.uk/testportal/job/api/notification'
+    declare -r API_ENDPOINT='http://www.adbraun.lol:9001'
+    jq -c . $VM_JSON
     curl \
     -X POST \
     -H "content-type: application/json" \
     -H "accept: application/json" \
     -H "Authorization: Token $API_TOKEN" \
-    -d '{ "job_id":"'$RD_JOB_ID'","execution_id":"'$RD_EXEC_ID'", "status":"ok", "data":{"name":"'$NAME'", "ip":"'$IPv4'", "uuid":"'$UUID'", "type":"'$TYPE'"}}' \
-    https://projects.eidf.ac.uk/testportal/job/api/notification
+    -d "'$(jq -c . $VM_JSON)'" \
+    $API_ENDPOINT
 }
 
 main() {
-    # Get list of VMs within project
-    local TABLE=$(openstack --os-cloud $OS_PROJECT server list | tail -n +4 | head -n -1 | sed 's/|//g')
+    initialize_json
+    while read -r VM; do
+        declare vm_uuid=$(jq -r '.ID' <<< "$VM")
+        declare vm_name=$(jq -r '.Name' <<< "$VM")
+        declare vm_ip=$(jq -r '.IP' <<< "$VM")
 
-    echo "$TABLE" | while read -r line; do
-        # Pattern to match user VMs
-        if echo "$line" | grep -qE "$OS_PROJECT\-$WIN_PROJECT\-wb|app|gpu"; then
-            post_request "$line" "CLIENT"
+        if [[ "$vm_name" =~ ^$OS_PROJECT-$WIN_PROJECT-dc01$ ]]; then
+            add_vm $vm_name $vm_uuid $vm_ip 'DC'
 
-        # Pattern to match DC01
-        elif echo "$line" | grep -qE "$OS_PROJECT\-$WIN_PROJECT\-dc01"; then
-            post_request "$line" "DC"
+        elif [[ "$vm_name" =~ ^$OS_PROJECT-$WIN_PROJECT-wb[0-9]+$ ]]; then
+            add_vm $vm_name $vm_uuid $vm_ip 'WB'
 
-        else
-            continue
+        elif [[ "$vm_name" =~ ^$OS_PROJECT-$WIN_PROJECT-app[0-9]+$ ]]; then
+            add_vm $vm_name $vm_uuid $vm_ip 'APP'
+
+        elif [[ "$vm_name" =~ ^$OS_PROJECT-$WIN_PROJECT-gpu[0-9]+$ ]]; then
+            add_vm $vm_name $vm_uuid $vm_ip 'GPU'
         fi
-    done
+    done < <(openstack --os-cloud $OS_PROJECT server list -f json | jq -c '.[] | {ID, Name, IP: .Networks.external[0]}')
+
+    post_request
 }
 
 main