From d18e86a9bdfe57f4999178e8f1e543477e14eae4 Mon Sep 17 00:00:00 2001
From: Artur Reznikov <artur.reznikov@avalabs.org>
Date: Thu, 27 Jun 2024 11:15:05 -0700
Subject: [PATCH] add security key

---
 examples/aws.go | 75 ++++++++++++++++++++++++++++++++++++++++++++++---
 utils/net.go    | 41 ++++++++++++++++++++++++++-
 2 files changed, 111 insertions(+), 5 deletions(-)

diff --git a/examples/aws.go b/examples/aws.go
index 1f96585..808b3b3 100644
--- a/examples/aws.go
+++ b/examples/aws.go
@@ -1,7 +1,7 @@
 // Copyright (C) 2024, Ava Labs, Inc. All rights reserved.
 // See the file LICENSE for licensing terms.
 
-package examples
+package main
 
 import (
 	"context"
@@ -9,21 +9,87 @@ import (
 	"time"
 
 	"github.com/ava-labs/avalanche-tooling-sdk-go/avalanche"
+	awsAPI "github.com/ava-labs/avalanche-tooling-sdk-go/cloud/aws"
 	"github.com/ava-labs/avalanche-tooling-sdk-go/node"
+	"github.com/ava-labs/avalanche-tooling-sdk-go/utils"
 )
 
+// createSecurityGroup creates a new security group in AWS using the specified AWS profile and region.
+//
+// ctx: The context.Context object for the request.
+// awsProfile: The AWS profile to use for the request.
+// awsRegion: The AWS region to use for the request.
+// Returns the ID of the created security group and an error, if any.
+func createSecurityGroup(ctx context.Context, awsProfile string, awsRegion string) (string, error) {
+	ec2Svc, err := awsAPI.NewAwsCloud(
+		ctx,
+		awsProfile,
+		awsRegion,
+	)
+	if err != nil {
+		return "", err
+	}
+	//detect user IP address
+	userIPAddress, err := utils.GetUserIPAddress()
+	if err != nil {
+		return "", err
+	}
+	const securityGroupName = "avalanche-tooling-sdk-go"
+	return ec2Svc.SetupSecurityGroup(userIPAddress, securityGroupName)
+}
+
+// createSSHKeyPair creates an SSH key pair for AWS.
+//
+// ctx: The context for the request.
+// awsProfile: The AWS profile to use for the request.
+// awsRegion: The AWS region to use for the request.
+// sshPrivateKeyPath: The path to save the SSH private key.
+// Returns an error if unable to create the key pair.
+func createSSHKeyPair(ctx context.Context, awsProfile string, awsRegion string, keyPairName string, sshPrivateKeyPath string) error {
+	if utils.FileExists(sshPrivateKeyPath) {
+		return fmt.Errorf("ssh private key %s already exists", sshPrivateKeyPath)
+	}
+	ec2Svc, err := awsAPI.NewAwsCloud(
+		ctx,
+		awsProfile,
+		awsRegion,
+	)
+	if err != nil {
+		return err
+	}
+	return ec2Svc.CreateAndDownloadKeyPair(keyPairName, sshPrivateKeyPath)
+}
+
 func main() {
 	ctx := context.Background()
+
 	// Get the default cloud parameters for AWS
 	cp, err := node.GetDefaultCloudParams(ctx, node.AWSCloud)
+	if err != nil {
+		panic(err)
+	}
 	// Set the cloud parameters for AWS non provided by the default
 	// Please set your own values for the following fields
-	cp.AWSConfig.AWSProfile = "default"
-	cp.AWSConfig.AWSSecurityGroupID = "sg-0e198c427f8f0616b"
-	cp.AWSConfig.AWSKeyPair = "artur-us-east-1-avalanche-cli"
+	// For example cp.AWSConfig.AWSProfile = "default" and so on
+
+	// Create security group or you can provide your own by setting cp.AWSConfig.AWSSecurityGroupID
+	sgID, err := createSecurityGroup(ctx, cp.AWSConfig.AWSProfile, cp.Region)
 	if err != nil {
 		panic(err)
 	}
+	cp.AWSConfig.AWSSecurityGroupID = sgID
+
+	// Change to your own path
+	sshPrivateKey := utils.ExpandHome("~/.ssh/avalanche-tooling-sdk-go.pem")
+
+	// Create aws ssh key pair or you can provide your own by setting cp.AWSConfig.AWSKeyPair.
+	// Make sure that ssh private key matches AWSKeyPair
+	keyPairName := "avalanche-tooling-sdk-go"
+	if err := createSSHKeyPair(ctx, cp.AWSConfig.AWSProfile, cp.Region, keyPairName, sshPrivateKey); err != nil {
+		panic(err)
+	}
+	cp.AWSConfig.AWSKeyPair = keyPairName
+
 	// Create a new host instance. Count is 1 so only one host will be created
 	const (
 		avalanchegoVersion  = "v1.11.8"
@@ -38,6 +104,7 @@ func main() {
 			AvalancheGoVersion:  avalanchegoVersion,
 			AvalancheCliVersion: avalancheCliVersion,
 			UseStaticIP:         false,
+			SSHPrivateKey:       sshPrivateKey,
 		})
 	if err != nil {
 		panic(err)
diff --git a/utils/net.go b/utils/net.go
index d82f925..98bc889 100644
--- a/utils/net.go
+++ b/utils/net.go
@@ -3,8 +3,47 @@
 
 package utils
 
-import "net"
+import (
+	"encoding/json"
+	"errors"
+	"io"
+	"net"
+	"net/http"
+)
 
 func IsValidIP(ipStr string) bool {
 	return net.ParseIP(ipStr) != nil
 }
+
+// GetUserIPAddress retrieves the IP address of the user.
+func GetUserIPAddress() (string, error) {
+	resp, err := http.Get("https://api.ipify.org?format=json")
+	if err != nil {
+		return "", err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return "", errors.New("HTTP request failed")
+	}
+
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return "", err
+	}
+
+	var result map[string]interface{}
+	if err := json.Unmarshal(body, &result); err != nil {
+		return "", err
+	}
+
+	ipAddress, ok := result["ip"].(string)
+	if ok {
+		if net.ParseIP(ipAddress) == nil {
+			return "", errors.New("invalid IP address")
+		}
+		return ipAddress, nil
+	}
+
+	return "", errors.New("no IP address found")
+}