populate CF Access headers if env is set (#111)

* create internal http/client package

* use http client in cmd package

* pass context everywhere

* mod download instead of tidy

* lint

* Update agent.go

Author: imperfect-fourth

Dockerfile

@@ -29,7 +29,7 @@ WORKDIR /app
 COPY go.mod go.sum ./
 
 # Download dependencies
-RUN go mod tidy
+RUN go mod download
 
 # Copy the source code
 COPY . .

cmd/agent.go

@@ -6,22 +6,27 @@ package cmd
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
+	"log"
+	"net/http"
+	"os"
+	"time"
+
 	"github.com/google/uuid"
+	"github.com/shinobistack/gokakashi/internal/http/client"
 	"github.com/shinobistack/gokakashi/internal/restapi/v1/agents"
 	"github.com/shinobistack/gokakashi/internal/restapi/v1/agenttasks"
 	"github.com/shinobistack/gokakashi/internal/restapi/v1/integrations"
 	"github.com/shinobistack/gokakashi/internal/restapi/v1/scans"
 	"github.com/shinobistack/gokakashi/pkg/registry/v1"
 	"github.com/shinobistack/gokakashi/pkg/scanner/v1"
 	"github.com/spf13/cobra"
-	"log"
-	"net/http"
-	"os"
-	"time"
 )
 
+type httpClientKey struct{}
+
 var agentCmd = &cobra.Command{
 	Use:   "agent",
 	Short: "Manage agents for GoKakashi",
@@ -30,7 +35,31 @@ var agentCmd = &cobra.Command{
 var agentStartCmd = &cobra.Command{
 	Use:   "start",
 	Short: "Register an agent and start polling for tasks",
-	Run:   agentRegister,
+	PersistentPreRun: func(cmd *cobra.Command, args []string) {
+		if token == "" {
+			log.Fatalf("Error: missing required flag --token")
+		}
+		headers := make(map[string]string)
+		cfClientID := os.Getenv("CF_ACCESS_CLIENT_ID")
+		cfClientSecret := os.Getenv("CF_ACCESS_CLIENT_SECRET")
+		if cfClientID != "" && cfClientSecret != "" {
+			headers["CF-Access-Client-Id"] = cfClientID
+			headers["CF-Access-Client-Secret"] = cfClientSecret
+		} else if cfClientSecret != "" {
+			fmt.Println("Warning: ignoring CF_ACCESS_CLIENT_SECRET because CF_ACCESS_CLIENT_ID is not set")
+		} else if cfClientID != "" {
+			fmt.Println("Warning: ignoring CF_ACCESS_CLIENT_ID because CF_ACCESS_CLIENT_SECRET is not set")
+		}
+
+		httpClient := client.New(
+			client.WithToken(token),
+			client.WithHeaders(headers),
+		)
+
+		ctx := context.WithValue(context.Background(), httpClientKey{}, httpClient)
+		cmd.SetContext(ctx)
+	},
+	Run: agentRegister,
 }
 
 var (
@@ -51,18 +80,18 @@ func agentRegister(cmd *cobra.Command, args []string) {
 	// log.Printf("Server: %s, Token: %s, Workspace: %s", server, token, workspace)
 
 	// Register the agent
-	agentID, err := registerAgent(server, token, workspace, name)
+	agentID, err := registerAgent(cmd.Context(), server, token, workspace, name)
 	if err != nil {
 		log.Fatalf("Failed to register the agent: %v", err)
 	}
 
 	log.Printf("Agent registered successfully! Agent ID: %d", agentID)
 
 	// Start polling for tasks
-	pollTasks(server, token, agentID, workspace)
+	pollTasks(cmd.Context(), server, token, agentID, workspace)
 }
 
-func registerAgent(server, token, workspace, name string) (int, error) {
+func registerAgent(ctx context.Context, server, token, workspace, name string) (int, error) {
 	reqBody := agents.RegisterAgentRequest{
 		Server:    server,
 		Token:     token,
@@ -75,10 +104,9 @@ func registerAgent(server, token, workspace, name string) (int, error) {
 	if err != nil {
 		return 0, fmt.Errorf("failed to create registration request: %w", err)
 	}
-	req.Header.Set("Authorization", "Bearer "+token)
 	req.Header.Set("Content-Type", "application/json")
 
-	resp, err := http.DefaultClient.Do(req)
+	resp, err := ctx.Value(httpClientKey{}).(*client.Client).Do(req)
 	if err != nil {
 		return 0, fmt.Errorf("failed to send registration request: %w", err)
 	}
@@ -95,10 +123,10 @@ func registerAgent(server, token, workspace, name string) (int, error) {
 	return response.ID, nil
 }
 
-func pollTasks(server, token string, agentID int, workspace string) {
+func pollTasks(ctx context.Context, server, token string, agentID int, workspace string) {
 	for {
 		// Process only tasks with status "pending" in the order returned (created_at ASC)
-		tasks, err := fetchTasks(server, token, agentID, "pending")
+		tasks, err := fetchTasks(ctx, server, token, agentID, "pending")
 		if err != nil {
 			log.Printf("Error fetching tasks: %v", err)
 			time.Sleep(10 * time.Second)
@@ -113,12 +141,12 @@ func pollTasks(server, token string, agentID int, workspace string) {
 
 		for _, task := range tasks {
 			// Update task status to "in_progress"
-			err := updateAgentTaskStatus(server, token, task.ID, agentID, "in_progress")
+			err := updateAgentTaskStatus(ctx, server, token, task.ID, agentID, "in_progress")
 			if err != nil {
 				log.Printf("Failed to update agent_task status to 'in_progress': %v", err)
 				return
 			}
-			processTask(server, token, task, workspace, agentID)
+			processTask(ctx, server, token, task, workspace, agentID)
 			continue
 		}
 		// Todo: Polling interval time decide
@@ -127,15 +155,13 @@ func pollTasks(server, token string, agentID int, workspace string) {
 	}
 }
 
-func fetchTasks(server, token string, agentID int, status string) ([]agenttasks.GetAgentTaskResponse, error) {
+func fetchTasks(ctx context.Context, server, token string, agentID int, status string) ([]agenttasks.GetAgentTaskResponse, error) {
 	req, err := http.NewRequest("GET", fmt.Sprintf("%s/api/v1/agents/%d/tasks?status=%s", server, agentID, status), nil)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create task polling request: %w", err)
 	}
 
-	req.Header.Set("Authorization", "Bearer "+token)
-
-	resp, err := http.DefaultClient.Do(req)
+	resp, err := ctx.Value(httpClientKey{}).(*client.Client).Do(req)
 	if err != nil {
 		return nil, fmt.Errorf("failed to send task polling request: %w", err)
 	}
@@ -153,7 +179,7 @@ func fetchTasks(server, token string, agentID int, status string) ([]agenttasks.
 	return tasks, nil
 }
 
-func updateAgentTaskStatus(server, token string, taskID uuid.UUID, agentID int, status string) error {
+func updateAgentTaskStatus(ctx context.Context, server, token string, taskID uuid.UUID, agentID int, status string) error {
 	reqBody := agenttasks.UpdateAgentTaskRequest{
 		ID:      taskID,
 		AgentID: intPtr(agentID),
@@ -166,10 +192,9 @@ func updateAgentTaskStatus(server, token string, taskID uuid.UUID, agentID int,
 	if err != nil {
 		return fmt.Errorf("failed to create task status update request: %w", err)
 	}
-	req.Header.Set("Authorization", "Bearer "+token)
 	req.Header.Set("Content-Type", "application/json")
 
-	resp, err := http.DefaultClient.Do(req)
+	resp, err := ctx.Value(httpClientKey{}).(*client.Client).Do(req)
 	if err != nil {
 		return fmt.Errorf("failed to update task status: %w", err)
 	}
@@ -182,16 +207,16 @@ func updateAgentTaskStatus(server, token string, taskID uuid.UUID, agentID int,
 	return nil
 }
 
-func processTask(server, token string, task agenttasks.GetAgentTaskResponse, workspace string, agentID int) {
+func processTask(ctx context.Context, server, token string, task agenttasks.GetAgentTaskResponse, workspace string, agentID int) {
 	// Step 1: Fetch scan details
-	scan, err := fetchScan(server, token, task.ScanID)
+	scan, err := fetchScan(ctx, server, token, task.ScanID)
 	if err != nil {
 		log.Printf("Failed to fetch scan details: %v", err)
 		return
 	}
 
 	// Step 2: Fetch integration details
-	integration, err := fetchIntegration(server, token, scan.IntegrationID)
+	integration, err := fetchIntegration(ctx, server, token, scan.IntegrationID)
 	if err != nil {
 		log.Printf("Failed to fetch integration details: %v", err)
 		return
@@ -203,7 +228,7 @@ func processTask(server, token string, task agenttasks.GetAgentTaskResponse, wor
 		return
 	}
 
-	err = updateScanStatus(server, token, scan.ID, "scan_in_progress")
+	err = updateScanStatus(ctx, server, token, scan.ID, "scan_in_progress")
 	if err != nil {
 		log.Printf("Failed to update scan status to 'scan_in_progress': %v", err)
 	}
@@ -212,16 +237,16 @@ func processTask(server, token string, task agenttasks.GetAgentTaskResponse, wor
 	reportPath, err := performScan(scan.Image, scan.Scanner)
 	if err != nil {
 		log.Printf("Failed to perform scan: %v", err)
-		if err := updateScanStatus(server, token, scan.ID, "error"); err != nil {
+		if err := updateScanStatus(ctx, server, token, scan.ID, "error"); err != nil {
 			log.Printf("Failed to update scan status to 'error': %v", err)
 		}
 		return
 	}
 
 	// Step 5: Upload the scan report
-	if err := uploadReport(server, token, scan.ID, reportPath); err != nil {
+	if err := uploadReport(ctx, server, token, scan.ID, reportPath); err != nil {
 		log.Printf("Failed to upload scan report: %v", err)
-		if err := updateScanStatus(server, token, scan.ID, "error"); err != nil {
+		if err := updateScanStatus(ctx, server, token, scan.ID, "error"); err != nil {
 			log.Printf("Failed to update scan status to 'error': %v", err)
 		}
 		return
@@ -231,24 +256,24 @@ func processTask(server, token string, task agenttasks.GetAgentTaskResponse, wor
 	// Todo: if exists update the status to notify_pending else complete
 	if scan.Notify == nil || len(*scan.Notify) == 0 {
 		log.Printf("No notify specified for scan ID: %s", scan.ID)
-		if err := updateScanStatus(server, token, scan.ID, "success"); err != nil {
+		if err := updateScanStatus(ctx, server, token, scan.ID, "success"); err != nil {
 			log.Printf("Failed to update scan status to 'success': %v", err)
 		}
 	} else {
-		err = updateScanStatus(server, token, scan.ID, "notify_pending")
+		err = updateScanStatus(ctx, server, token, scan.ID, "notify_pending")
 		if err != nil {
 			log.Printf("Failed to update scan status to 'scan_in_progress': %v", err)
 		}
 	}
 
-	if err := updateAgentTaskStatus(server, token, task.ID, agentID, "complete"); err != nil {
+	if err := updateAgentTaskStatus(ctx, server, token, task.ID, agentID, "complete"); err != nil {
 		log.Printf("Failed to update agent_task status to 'complete': %v", err)
 	}
 
 	log.Printf("AgentTaskID completed successfully: %v", task.ID)
 }
 
-func updateScanStatus(server, token string, scanID uuid.UUID, status string) error {
+func updateScanStatus(ctx context.Context, server, token string, scanID uuid.UUID, status string) error {
 	reqBody := scans.UpdateScanRequest{
 		ID:     scanID,
 		Status: strPtr(status),
@@ -259,10 +284,9 @@ func updateScanStatus(server, token string, scanID uuid.UUID, status string) err
 	if err != nil {
 		return fmt.Errorf("failed to create scan status update request: %w", err)
 	}
-	req.Header.Set("Authorization", "Bearer "+token)
 	req.Header.Set("Content-Type", "application/json")
 
-	resp, err := http.DefaultClient.Do(req)
+	resp, err := ctx.Value(httpClientKey{}).(*client.Client).Do(req)
 	if err != nil {
 		return fmt.Errorf("failed to update scan status: %w", err)
 	}
@@ -275,14 +299,13 @@ func updateScanStatus(server, token string, scanID uuid.UUID, status string) err
 	return nil
 }
 
-func fetchScan(server, token string, scanID uuid.UUID) (*scans.GetScanResponse, error) {
+func fetchScan(ctx context.Context, server, token string, scanID uuid.UUID) (*scans.GetScanResponse, error) {
 	req, err := http.NewRequest("GET", fmt.Sprintf("%s/api/v1/scans/%s", server, scanID), nil)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create scan request: %w", err)
 	}
-	req.Header.Set("Authorization", "Bearer "+token)
 
-	resp, err := http.DefaultClient.Do(req)
+	resp, err := ctx.Value(httpClientKey{}).(*client.Client).Do(req)
 	if err != nil {
 		return nil, fmt.Errorf("failed to fetch scan details: %w", err)
 	}
@@ -300,14 +323,13 @@ func fetchScan(server, token string, scanID uuid.UUID) (*scans.GetScanResponse,
 	return &scan, nil
 }
 
-func fetchIntegration(server, token string, integrationID uuid.UUID) (*integrations.GetIntegrationResponse, error) {
+func fetchIntegration(ctx context.Context, server, token string, integrationID uuid.UUID) (*integrations.GetIntegrationResponse, error) {
 	req, err := http.NewRequest("GET", fmt.Sprintf("%s/api/v1/integrations/%s", server, integrationID), nil)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create integration fetch request: %w", err)
 	}
-	req.Header.Set("Authorization", "Bearer "+token)
 
-	resp, err := http.DefaultClient.Do(req)
+	resp, err := ctx.Value(httpClientKey{}).(*client.Client).Do(req)
 	if err != nil {
 		return nil, fmt.Errorf("failed to fetch integration details: %w", err)
 	}
@@ -363,7 +385,7 @@ func performScan(image, scannerType string) (string, error) {
 	return reportPath, nil
 }
 
-func uploadReport(server, token string, scanID uuid.UUID, reportPath string) error {
+func uploadReport(ctx context.Context, server, token string, scanID uuid.UUID, reportPath string) error {
 	report, err := os.ReadFile(reportPath)
 	if err != nil {
 		return fmt.Errorf("failed to read report file: %w", err)
@@ -383,10 +405,9 @@ func uploadReport(server, token string, scanID uuid.UUID, reportPath string) err
 		return fmt.Errorf("failed to create report upload request: %w", err)
 	}
 
-	req.Header.Set("Authorization", "Bearer "+token)
 	req.Header.Set("Content-Type", "application/json")
 
-	resp, err := http.DefaultClient.Do(req)
+	resp, err := ctx.Value(httpClientKey{}).(*client.Client).Do(req)
 	if err != nil {
 		return fmt.Errorf("failed to upload scan report: %w", err)
 	}

internal/http/client/client.go

@@ -0,0 +1,36 @@
+package client
+
+import (
+	"fmt"
+	"net/http"
+)
+
+type Client struct {
+	client  *http.Client
+	token   string
+	headers map[string]string
+}
+
+func New(opts ...Option) *Client {
+	c := &Client{
+		client: http.DefaultClient,
+	}
+
+	for _, o := range opts {
+		o(c)
+	}
+
+	return c
+}
+
+func (c *Client) Do(req *http.Request) (*http.Response, error) {
+	if c.token != "" {
+		req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", c.token))
+	}
+
+	for k, v := range c.headers {
+		req.Header.Set(k, v)
+	}
+
+	return c.client.Do(req)
+}

internal/http/client/client_opts.go

@@ -0,0 +1,23 @@
+package client
+
+import "net/http"
+
+type Option func(*Client)
+
+func WithHTTPClient(client *http.Client) Option {
+	return func(c *Client) {
+		c.client = client
+	}
+}
+
+func WithToken(token string) Option {
+	return func(c *Client) {
+		c.token = token
+	}
+}
+
+func WithHeaders(headers map[string]string) Option {
+	return func(c *Client) {
+		c.headers = headers
+	}
+}