From 91c56221b74d0c7087d912462153805e6545f75b Mon Sep 17 00:00:00 2001
From: zhewang <zwan@redhat.com>
Date: Thu, 12 Dec 2024 16:59:51 +0800
Subject: [PATCH] OCM-13040 | test: Bastion proxy support username and password

---
 pkg/test/vpc_client/bastion.go | 106 +++++++++++++++++++--------------
 1 file changed, 62 insertions(+), 44 deletions(-)

diff --git a/pkg/test/vpc_client/bastion.go b/pkg/test/vpc_client/bastion.go
index 913ace6..094928b 100644
--- a/pkg/test/vpc_client/bastion.go
+++ b/pkg/test/vpc_client/bastion.go
@@ -6,11 +6,13 @@ import (
 	"fmt"
 	"github.com/aws/aws-sdk-go-v2/service/ec2/types"
 	"github.com/openshift-online/ocm-common/pkg/file"
-	"net"
+	"golang.org/x/crypto/bcrypt"
+	"path"
 	"time"
 
 	CON "github.com/openshift-online/ocm-common/pkg/aws/consts"
 	"github.com/openshift-online/ocm-common/pkg/log"
+	"github.com/openshift-online/ocm-common/pkg/utils"
 )
 
 // LaunchBastion will launch a bastion instance on the indicated zone.
@@ -94,55 +96,45 @@ func (vpc *VPC) LaunchBastion(imageID string, zone string, userData string, keyp
 	return inst, nil
 }
 
-func (vpc *VPC) PrepareBastionProxy(zone string, cidrBlock string, keypairName string,
-	privateKeyPath string) (*types.Instance, error) {
-	filters := []map[string][]string{
-		{
-			"vpc-id": {
-				vpc.VpcID,
-			},
-		},
-		{
-			"tag:Name": {
-				CON.BastionName,
-			},
-		},
-	}
-
-	insts, err := vpc.AWSClient.ListInstances([]string{}, filters...)
+// PrepareBastionProxy will launch a bastion instance with squid proxy on the indicated zone and return the proxy url.
+func (vpc *VPC) PrepareBastionProxy(zone string, keypairName string, privateKeyPath string) (proxyUrl string, err error) {
+	userData := generateShellCommand()
+
+	encodeUserData := base64.StdEncoding.EncodeToString([]byte(userData))
+	instance, err := vpc.LaunchBastion("", zone, encodeUserData, keypairName, privateKeyPath)
 	if err != nil {
-		return nil, err
+		log.LogError("Launch bastion failed")
 	}
-	if len(insts) == 0 {
-		log.LogInfo("Didn't found an existing bastion, going to launch one")
-		if cidrBlock == "" {
-			cidrBlock = CON.RouteDestinationCidrBlock
-		}
-		_, _, err = net.ParseCIDR(cidrBlock)
-		if err != nil {
-			log.LogError("CIDR IP address format is invalid")
-			return nil, err
-		}
 
-		userData := fmt.Sprintf(`#!/bin/bash
-		yum update -y
-		yum install -y squid
-		cd /etc/squid/
-		sudo mv ./squid.conf ./squid.conf.bak
-		sudo touch squid.conf
-		echo http_port 3128 >> /etc/squid/squid.conf
-		echo acl allowed_ips src %s >> /etc/squid/squid.conf
-		echo http_access allow allowed_ips >> /etc/squid/squid.conf
-		echo http_access deny all >> /etc/squid/squid.conf
-		systemctl start squid
-		systemctl enable squid`, cidrBlock)
+	username := utils.RandomLabel(5)
+	password := utils.RandomLabel(10)
+
+	hashedPassword, err := generateBcryptPassword(password)
+	if err != nil {
+		return "", err
+	}
 
-		encodeUserData := base64.StdEncoding.EncodeToString([]byte(userData))
-		return vpc.LaunchBastion("", zone, encodeUserData, keypairName, privateKeyPath)
+	line := fmt.Sprintf("%s:%s\n", username, hashedPassword)
+	remoteFilePath := "/etc/squid/passwords"
+	hostName := fmt.Sprintf("%s:22", *instance.PublicIpAddress)
+	createFileCMD := fmt.Sprintf("sudo touch %s", remoteFilePath)
+	copyPasswordCMD := fmt.Sprintf("echo '%s' | sudo tee %s > /dev/null", line, remoteFilePath)
+	SSHExecuteCMDs := []string{
+		createFileCMD,
+		copyPasswordCMD,
+	}
+	privateKeyName := fmt.Sprintf("%s-%s", path.Join(privateKeyPath, keypairName), "keyPair.pem")
 
+	for _, cmd := range SSHExecuteCMDs {
+		_, err = Exec_CMD(CON.AWSInstanceUser, privateKeyName, hostName, cmd)
+		if err != nil {
+			log.LogError("SSH execute command failed")
+			return "", err
+		}
 	}
-	log.LogInfo("Found existing bastion: %s", *insts[0].InstanceId)
-	return &insts[0], nil
+
+	proxyUrl = fmt.Sprintf("http://%s:%s@%s:3128", username, password, *instance.PublicIpAddress)
+	return proxyUrl, nil
 }
 
 func (vpc *VPC) DestroyBastionProxy(instance types.Instance) error {
@@ -155,3 +147,29 @@ func (vpc *VPC) DestroyBastionProxy(instance types.Instance) error {
 	}
 	return nil
 }
+
+func generateBcryptPassword(plainPassword string) (string, error) {
+	hashedPassword, err := bcrypt.GenerateFromPassword([]byte(plainPassword), bcrypt.DefaultCost)
+	if err != nil {
+		log.LogError("Generate hashed password failed")
+		return "", nil
+	}
+	return string(hashedPassword), nil
+}
+
+func generateShellCommand() string {
+	return `#!/bin/bash
+		yum update -y
+		sudo dnf install squid -y
+		cd /etc/squid/
+		sudo mv ./squid.conf ./squid.conf.bak
+		sudo touch squid.conf
+		echo http_port 3128 >> /etc/squid/squid.conf
+		echo auth_param basic program /usr/lib64/squid/basic_ncsa_auth /etc/squid/passwords >> /etc/squid/squid.conf
+		echo auth_param basic realm Squid Proxy Server >> /etc/squid/squid.conf
+		echo acl authenticated proxy_auth REQUIRED >> /etc/squid/squid.conf
+		echo http_access allow authenticated >> /etc/squid/squid.conf
+		echo http_access deny all >> /etc/squid/squid.conf
+		systemctl start squid
+		systemctl enable squid`
+}