From e8962c3cba8c47c73b7bce93466d7bd6e44231fe Mon Sep 17 00:00:00 2001
From: schulze <thefeli73@hotmail.com>
Date: Mon, 17 Apr 2023 15:54:42 +0200
Subject: [PATCH] Availability zone randomisation

---
 mtdaws/mtd.go   |  6 ++++--
 mtdaws/utils.go | 54 ++++++++++++++++++++++++++++++++++++++++++++++---
 2 files changed, 55 insertions(+), 5 deletions(-)

diff --git a/mtdaws/mtd.go b/mtdaws/mtd.go
index 3f81795..8e428ae 100644
--- a/mtdaws/mtd.go
+++ b/mtdaws/mtd.go
@@ -7,6 +7,7 @@ import (
 
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/ec2"
+	"github.com/google/uuid"
 	"github.com/thefeli73/polemos/state"
 )
 
@@ -38,6 +39,7 @@ func AWSUpdateService(config state.Config, region string, service state.CustomUU
 
 // AWSMoveInstance moves a specified instance to a new availability region
 func AWSMoveInstance(config state.Config) (state.Config) {
+
 	// pseudorandom instance from all services for testing
 	var serviceUUID state.CustomUUID
 	var instance state.Service
@@ -47,7 +49,7 @@ func AWSMoveInstance(config state.Config) (state.Config) {
 		break
 	}
 
-	fmt.Println("MTD move service:\t", serviceUUID)
+	fmt.Println("MTD move service:\t", uuid.UUID.String(uuid.UUID(serviceUUID)))
 
 	region, instanceID := DecodeCloudID(instance.CloudID)
 	awsConfig := NewConfig(region, config.AWS.CredentialsPath)
@@ -73,7 +75,7 @@ func AWSMoveInstance(config state.Config) (state.Config) {
 	}
 	fmt.Println("Image is ready:\t\t", imageName)
 
-	newInstanceID, err := launchInstance(svc, realInstance, imageName)
+	newInstanceID, err := launchInstance(svc, realInstance, imageName, region)
 	if err != nil {
 		fmt.Println("Error launching instance:\t", err)
 		return config
diff --git a/mtdaws/utils.go b/mtdaws/utils.go
index 32e87ed..2d4d5a4 100644
--- a/mtdaws/utils.go
+++ b/mtdaws/utils.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"math/rand"
 	"os"
 	"strings"
 	"time"
@@ -141,14 +142,18 @@ func waitForImageReady(svc *ec2.Client, imageID string, timeout time.Duration) e
 	}
 }
 
-// launchInstance launches a instance based on an oldInstance and AMI (duplicating the instance)
-func launchInstance(svc *ec2.Client, oldInstance *types.Instance, imageID string) (string, error) {
+// launchInstance launches a instance IN RANDOM AVAILABILITY ZONE within the same region, based on an oldInstance and AMI (duplicating the instance)
+func launchInstance(svc *ec2.Client, oldInstance *types.Instance, imageID string, region string) (string, error) {
 	securityGroupIds := make([]string, len(oldInstance.SecurityGroups))
 	for i, sg := range oldInstance.SecurityGroups {
 		securityGroupIds[i] = aws.ToString(sg.GroupId)
 	}
 	// TODO: select random zone that is not the current one.
-	availabilityZone := "us-east-1d"
+	availabilityZone, err := getRandomDifferentAvailabilityZone(svc, oldInstance, region)
+	if err != nil {
+		return "", err
+	}
+
 	input := &ec2.RunInstancesInput{
 		ImageId:         aws.String(imageID),
 		InstanceType:    oldInstance.InstanceType,
@@ -171,6 +176,49 @@ func launchInstance(svc *ec2.Client, oldInstance *types.Instance, imageID string
 	return aws.ToString(output.Instances[0].InstanceId), nil
 }
 
+// getRandomDifferentAvailabilityZone fetches all AZ from the same region as the instance and returns a random AZ that is not equal to the one used by the instance
+func getRandomDifferentAvailabilityZone(svc *ec2.Client, instance *types.Instance, region string) (string, error) {
+	// Seed the random generator
+	rand.Seed(time.Now().UnixNano())
+
+	// Get the current availability zone of the instance
+	currentAZ := aws.ToString(instance.Placement.AvailabilityZone)
+
+	// Describe availability zones in the region
+	input := &ec2.DescribeAvailabilityZonesInput{
+		Filters: []types.Filter{
+			{
+				Name:   aws.String("region-name"),
+				Values: []string{region},
+			},
+		},
+	}
+
+	output, err := svc.DescribeAvailabilityZones(context.TODO(), input)
+	if err != nil {
+		return "", err
+	}
+
+	// Filter out the current availability zone
+	availableAZs := []string{}
+	for _, az := range output.AvailabilityZones {
+		if aws.ToString(az.ZoneName) != currentAZ {
+			availableAZs = append(availableAZs, aws.ToString(az.ZoneName))
+		}
+	}
+
+	// If no other availability zones are available, return an error
+	if len(availableAZs) == 0 {
+		return "", errors.New("no other availability zones available")
+	}
+
+	// Select a random availability zone from the remaining ones
+	randomIndex := rand.Intn(len(availableAZs))
+	randomAZ := availableAZs[randomIndex]
+	return randomAZ, nil
+}
+
+
 // terminateInstance kills an instance by id
 func terminateInstance(svc *ec2.Client, instanceID string) error {
 	input := &ec2.TerminateInstancesInput{