From 282745e0a755022f0573f4ef362b4d3d96ea25d1 Mon Sep 17 00:00:00 2001
From: Chaitanya Baraskar <chaitanya.axe@gmail.com>
Date: Wed, 7 Oct 2020 18:19:31 +0530
Subject: [PATCH] Fixed happy path flow now SG gets created in default VPC.
 Also fixed 389. (#670)

---
 cmd/instance.go |  2 +-
 lepton/aws.go   | 14 +++++++++++---
 lepton/utils.go | 11 +++++++++++
 3 files changed, 23 insertions(+), 4 deletions(-)
 create mode 100644 lepton/utils.go

diff --git a/cmd/instance.go b/cmd/instance.go
index 3f51fb92..b1ecf3cf 100644
--- a/cmd/instance.go
+++ b/cmd/instance.go
@@ -39,7 +39,7 @@ func instanceCreateCommandHandler(cmd *cobra.Command, args []string) {
 	}
 
 	zone, _ := cmd.Flags().GetString("zone")
-	if zone == "" && (provider == "gcp" || provider == "aws") {
+	if zone != "" && (provider == "gcp" || provider == "aws") {
 		c.CloudConfig.Zone = zone
 	}
 
diff --git a/lepton/aws.go b/lepton/aws.go
index 1b4d1adb..035d36f5 100644
--- a/lepton/aws.go
+++ b/lepton/aws.go
@@ -396,7 +396,7 @@ func (p *AWS) SyncImage(config *Config, target Provider, image string) error {
 func (p *AWS) CreateInstance(ctx *Context) error {
 	result, err := getAWSImages(ctx.config.CloudConfig.Zone)
 	if err != nil {
-		return err
+		exitWithError("Invalid zone")
 	}
 
 	imgName := ctx.config.CloudConfig.ImageName
@@ -493,7 +493,15 @@ func (p *AWS) CreateSG(ctx *Context, svc *ec2.EC2, imgName string) (string, erro
 	if len(result.Vpcs) == 0 {
 		fmt.Println("No VPCs found to associate security group with.")
 	}
-	vpcID := aws.StringValue(result.Vpcs[0].VpcId)
+
+	var vpcID string
+	//TODO: This will fail if there is no default VPC. Need to implement feature where user should be able to mention VPC
+	for i, s := range result.Vpcs {
+		isDefault := *s.IsDefault
+		if isDefault == true {
+			vpcID = aws.StringValue(result.Vpcs[i].VpcId)
+		}
+	}
 
 	t := time.Now().UnixNano()
 	s := strconv.FormatInt(t, 10)
@@ -525,7 +533,7 @@ func (p *AWS) CreateSG(ctx *Context, svc *ec2.EC2, imgName string) (string, erro
 
 	// maybe have these ports specified from config.json in near future
 	_, err = svc.AuthorizeSecurityGroupIngress(&ec2.AuthorizeSecurityGroupIngressInput{
-		GroupName: aws.String(sgName),
+		GroupId: createRes.GroupId,
 		IpPermissions: []*ec2.IpPermission{
 			(&ec2.IpPermission{}).
 				SetIpProtocol("tcp").
diff --git a/lepton/utils.go b/lepton/utils.go
new file mode 100644
index 00000000..4329435e
--- /dev/null
+++ b/lepton/utils.go
@@ -0,0 +1,11 @@
+package lepton
+
+import (
+	"fmt"
+	"os"
+)
+
+func exitWithError(errs string) {
+	fmt.Println(fmt.Sprintf(ErrorColor, errs))
+	os.Exit(1)
+}