From 913acdc2aec338ca293a2d5d37315f6598bc51d5 Mon Sep 17 00:00:00 2001
From: Zachary Winnerman <98712682+zwinnerman-fleetdm@users.noreply.github.com>
Date: Mon, 25 Jul 2022 16:26:04 -0400
Subject: [PATCH] Add auth header requirement for /new and fix openapi spec for
 timestamp (#6855)

---
 .../sandbox/JITProvisioner/jitprovisioner.tf  |  6 +++++
 .../sandbox/JITProvisioner/lambda/main.go     | 24 +++++++++++--------
 2 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/infrastructure/sandbox/JITProvisioner/jitprovisioner.tf b/infrastructure/sandbox/JITProvisioner/jitprovisioner.tf
index 24b14bac90..ae1cb3137c 100644
--- a/infrastructure/sandbox/JITProvisioner/jitprovisioner.tf
+++ b/infrastructure/sandbox/JITProvisioner/jitprovisioner.tf
@@ -130,10 +130,16 @@ resource "aws_lambda_function" "jitprovisioner" {
       DYNAMODB_LIFECYCLE_TABLE = var.dynamodb_table.id
       LIFECYCLE_SFN            = aws_sfn_state_machine.main.arn
       FLEET_BASE_URL           = "${var.base_domain}"
+      AUTHORIZATION_PSK        = random_password.authorization.result
     }
   }
 }
 
+resource "random_password" "authorization" {
+  length  = 16
+  special = false
+}
+
 output "jitprovisioner" {
   value = aws_lambda_function.jitprovisioner
 }
diff --git a/infrastructure/sandbox/JITProvisioner/lambda/main.go b/infrastructure/sandbox/JITProvisioner/lambda/main.go
index 5ad4f04a44..ee554c4bfc 100644
--- a/infrastructure/sandbox/JITProvisioner/lambda/main.go
+++ b/infrastructure/sandbox/JITProvisioner/lambda/main.go
@@ -7,6 +7,7 @@ import (
 	flags "github.com/jessevdk/go-flags"
 	//"github.com/juju/errors"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/arn"
@@ -24,7 +25,6 @@ import (
 	"math/rand"
 	"strings"
 	"time"
-	"errors"
 )
 
 type OptionsStruct struct {
@@ -32,6 +32,7 @@ type OptionsStruct struct {
 	LifecycleTable     string `long:"dynamodb-lifecycle-table" env:"DYNAMODB_LIFECYCLE_TABLE" required:"true"`
 	LifecycleSFN       string `long:"lifecycle-sfn" env:"LIFECYCLE_SFN" required:"true"`
 	FleetBaseURL       string `long:"fleet-base-url" env:"FLEET_BASE_URL" required:"true"`
+	AuthorizationPSK   string `long:"authorization-psk" env:"AUTHORIZATION_PSK" required:"true"`
 }
 
 var options = OptionsStruct{}
@@ -142,8 +143,8 @@ func triggerSFN(id, expiry string) (err error) {
 		return
 	}
 	if int(endTime.Sub(time.Now()).Seconds()) < 0 {
-	    return errors.New("Expiry time is in the past")
-    }
+		return errors.New("Expiry time is in the past")
+	}
 	sfnInStr, err := json.Marshal(struct {
 		InstanceID string `json:"instanceID"`
 		WaitTime   int    `json:"waitTime"`
@@ -180,12 +181,17 @@ type NewFleetInput struct {
 	Name              string `json:"name" validate:"required"`
 	SandboxExpiration string `json:"sandbox_expiration" validate:"required"`
 	Password          string `json:"password" validate:"required"`
+	Authorization     string `header:"Authorization" validate:"required"`
 }
 type NewFleetOutput struct {
 	URL string
 }
 
 func NewFleet(c *gin.Context, in *NewFleetInput) (ret *NewFleetOutput, err error) {
+	if in.Authorization != options.AuthorizationPSK {
+		err = errors.New("Unauthorized")
+		return
+	}
 	ret = &NewFleetOutput{}
 	fleet, err := getFleetInstance()
 	if err != nil {
@@ -202,9 +208,9 @@ func NewFleet(c *gin.Context, in *NewFleetInput) (ret *NewFleetOutput, err error
 	}
 	log.Print("Creating admin user")
 	if _, err = client.Setup(in.Email, in.Name, in.Password, fleet.ID); err != nil {
-	    log.Print(err)
-	    return
-    }
+		log.Print(err)
+		return
+	}
 	if err = triggerSFN(fleet.ID, in.SandboxExpiration); err != nil {
 		log.Print(err)
 		return
@@ -216,16 +222,14 @@ type ExpiryInput struct {
 	ID string `query:"id" validate:"required"`
 }
 type ExpiryOutput struct {
-    Timestamp string `json:"timestamp"`
+	Timestamp time.Time `json:"timestamp"`
 }
 
 func GetExpiry(c *gin.Context, in *ExpiryInput) (ret *ExpiryOutput, err error) {
 	ret = &ExpiryOutput{}
-	var expiry time.Time
-	if expiry, err = getExpiry(in.ID); err != nil {
+	if ret.Timestamp, err = getExpiry(in.ID); err != nil {
 		return
 	}
-	ret.Timestamp = expiry.Format(time.RFC3339)
 	return
 }