feature: skip image download if it exists

sztp-agent/cmd/daemon.go

@@ -59,7 +59,8 @@ func Daemon() *cobra.Command {
 					return fmt.Errorf("must not be folder: %q", filePath)
 				}
 			}
-			a := secureagent.NewAgent(bootstrapURL, serialNumber, dhcpLeaseFile, devicePassword, devicePrivateKey, deviceEndEntityCert, bootstrapTrustAnchorCert)
+			client := secureagent.NewHttpClient(bootstrapTrustAnchorCert, deviceEndEntityCert, devicePrivateKey)
+			a := secureagent.NewAgent(bootstrapURL, serialNumber, dhcpLeaseFile, devicePassword, devicePrivateKey, deviceEndEntityCert, bootstrapTrustAnchorCert, &client)
 			return a.RunCommandDaemon()
 		},
 	}

sztp-agent/cmd/disable.go

@@ -34,7 +34,8 @@ func Disable() *cobra.Command {
 		Use:   "disable",
 		Short: "Run the disable command",
 		RunE: func(_ *cobra.Command, _ []string) error {
-			a := secureagent.NewAgent(bootstrapURL, serialNumber, dhcpLeaseFile, devicePassword, devicePrivateKey, deviceEndEntityCert, bootstrapTrustAnchorCert)
+			client := secureagent.NewHttpClient(bootstrapTrustAnchorCert, deviceEndEntityCert, devicePrivateKey)
+			a := secureagent.NewAgent(bootstrapURL, serialNumber, dhcpLeaseFile, devicePassword, devicePrivateKey, deviceEndEntityCert, bootstrapTrustAnchorCert, &client)
 			return a.RunCommandDisable()
 		},
 	}

sztp-agent/cmd/enable.go

@@ -34,7 +34,8 @@ func Enable() *cobra.Command {
 		Use:   "enable",
 		Short: "Run the enable command",
 		RunE: func(_ *cobra.Command, _ []string) error {
-			a := secureagent.NewAgent(bootstrapURL, serialNumber, dhcpLeaseFile, devicePassword, devicePrivateKey, deviceEndEntityCert, bootstrapTrustAnchorCert)
+			client := secureagent.NewHttpClient(bootstrapTrustAnchorCert, deviceEndEntityCert, devicePrivateKey)
+			a := secureagent.NewAgent(bootstrapURL, serialNumber, dhcpLeaseFile, devicePassword, devicePrivateKey, deviceEndEntityCert, bootstrapTrustAnchorCert, &client)
 			return a.RunCommandEnable()
 		},
 	}

sztp-agent/cmd/run.go

@@ -59,7 +59,8 @@ func Run() *cobra.Command {
 					return fmt.Errorf("must not be folder: %q", filePath)
 				}
 			}
-			a := secureagent.NewAgent(bootstrapURL, serialNumber, dhcpLeaseFile, devicePassword, devicePrivateKey, deviceEndEntityCert, bootstrapTrustAnchorCert)
+			client := secureagent.NewHttpClient(bootstrapTrustAnchorCert, deviceEndEntityCert, devicePrivateKey)
+			a := secureagent.NewAgent(bootstrapURL, serialNumber, dhcpLeaseFile, devicePassword, devicePrivateKey, deviceEndEntityCert, bootstrapTrustAnchorCert, &client)
 			return a.RunCommand()
 		},
 	}

sztp-agent/cmd/status.go

@@ -34,7 +34,8 @@ func Status() *cobra.Command {
 		Use:   "status",
 		Short: "Run the status command",
 		RunE: func(_ *cobra.Command, _ []string) error {
-			a := secureagent.NewAgent(bootstrapURL, serialNumber, dhcpLeaseFile, devicePassword, devicePrivateKey, deviceEndEntityCert, bootstrapTrustAnchorCert)
+			client := secureagent.NewHttpClient(bootstrapTrustAnchorCert, deviceEndEntityCert, devicePrivateKey)
+			a := secureagent.NewAgent(bootstrapURL, serialNumber, dhcpLeaseFile, devicePassword, devicePrivateKey, deviceEndEntityCert, bootstrapTrustAnchorCert, &client)
 			return a.RunCommandStatus()
 		},
 	}

sztp-agent/pkg/secureagent/agent.go

@@ -9,6 +9,13 @@
 // Package secureagent implements the secure agent
 package secureagent
 
+import (
+	"crypto/tls"
+	"crypto/x509"
+	"net/http"
+	"os"
+)
+
 const (
 	CONTENT_TYPE_YANG = "application/yang-data+json"
 	OS_RELEASE_FILE   = "/etc/os-release"
@@ -68,6 +75,11 @@
 	} `json:"ietf-restconf:errors"`
 }
 
+type HttpGetter interface {
+	Get(uri string) (*http.Response, error)
+	Do(req *http.Request) (*http.Response, error)
+}
+
 // Agent is the basic structure to define an agent instance
 type Agent struct {
 	InputBootstrapURL             string                        // Bootstrap complete URL given by USER
@@ -83,10 +95,10 @@
 	ProgressJSON                  ProgressJSON                  // ProgressJson structure
 	BootstrapServerOnboardingInfo BootstrapServerOnboardingInfo // BootstrapServerOnboardingInfo structure
 	BootstrapServerRedirectInfo   BootstrapServerRedirectInfo   // BootstrapServerRedirectInfo structure
-
+	HttpClient                    HttpGetter
 }
 
-func NewAgent(bootstrapURL, serialNumber, dhcpLeaseFile, devicePassword, devicePrivateKey, deviceEndEntityCert, bootstrapTrustAnchorCert string) *Agent {
+func NewAgent(bootstrapURL, serialNumber, dhcpLeaseFile, devicePassword, devicePrivateKey, deviceEndEntityCert, bootstrapTrustAnchorCert string, httpClient HttpGetter) *Agent {
 	return &Agent{
 		InputBootstrapURL:             bootstrapURL,
 		BootstrapURL:                  "",
@@ -101,6 +113,7 @@
 		ProgressJSON:                  ProgressJSON{},
 		BootstrapServerRedirectInfo:   BootstrapServerRedirectInfo{},
 		BootstrapServerOnboardingInfo: BootstrapServerOnboardingInfo{},
+		HttpClient:                    httpClient,
 	}
 }
 
@@ -171,3 +184,25 @@
 func (a *Agent) SetProgressJSON(p ProgressJSON) {
 	a.ProgressJSON = p
 }
+
+func NewHttpClient(bootstrapTrustAnchorCert string, deviceEndEntityCert string, devicePrivateKey string) http.Client {
+	caCert, _ := os.ReadFile(bootstrapTrustAnchorCert)
+	caCertPool := x509.NewCertPool()
+	caCertPool.AppendCertsFromPEM(caCert)
+	cert, _ := tls.LoadX509KeyPair(deviceEndEntityCert, devicePrivateKey)
+	client := http.Client{
+		CheckRedirect: func(r *http.Request, _ []*http.Request) error {
+			r.URL.Opaque = r.URL.Path
+			return nil
+		},
+		Transport: &http.Transport{
+			TLSClientConfig: &tls.Config{
+				//nolint:gosec
+				InsecureSkipVerify: true, // TODO: remove skip verify
+				RootCAs:            caCertPool,
+				Certificates:       []tls.Certificate{cert},
+			},
+		},
+	}
+	return client
+}

sztp-agent/pkg/secureagent/agent_test.go

@@ -9,6 +9,7 @@ Copyright (C) 2022 Red Hat.
 package secureagent
 
 import (
+	"net/http"
 	"reflect"
 	"testing"
 )
@@ -829,6 +830,7 @@ func TestNewAgent(t *testing.T) {
 		deviceEndEntityCert      string
 		bootstrapTrustAnchorCert string
 	}
+	client := http.Client{}
 	tests := []struct {
 		name string
 		args args
@@ -856,12 +858,13 @@ func TestNewAgent(t *testing.T) {
 				ContentTypeReq:           "application/yang-data+json",
 				InputJSONContent:         generateInputJSONContent(),
 				DhcpLeaseFile:            "TestDhcpLeaseFile",
+				HttpClient:               &client,
 			},
 		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			if got := NewAgent(tt.args.bootstrapURL, tt.args.serialNumber, tt.args.dhcpLeaseFile, tt.args.devicePassword, tt.args.devicePrivateKey, tt.args.deviceEndEntityCert, tt.args.bootstrapTrustAnchorCert); !reflect.DeepEqual(got, tt.want) {
+			if got := NewAgent(tt.args.bootstrapURL, tt.args.serialNumber, tt.args.dhcpLeaseFile, tt.args.devicePassword, tt.args.devicePrivateKey, tt.args.deviceEndEntityCert, tt.args.bootstrapTrustAnchorCert, &client); !reflect.DeepEqual(got, tt.want) {
 				t.Errorf("NewAgent() = %v, want %v", got, tt.want)
 			}
 		})

sztp-agent/pkg/secureagent/daemon.go

@@ -10,16 +10,13 @@ package secureagent
 
 import (
 	"bytes"
-	"crypto/tls"
-	"crypto/x509"
 	"encoding/asn1"
 	"encoding/base64"
 	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
 	"log"
-	"net/http"
 	"net/url"
 	"os"
 	"os/exec"
@@ -181,88 +178,106 @@ func (a *Agent) doRequestBootstrapServerOnboardingInfo() error {
 	return errri
 }
 
-//nolint:funlen
-func (a *Agent) downloadAndValidateImage() error {
-	log.Printf("[INFO] Starting the Download Image: %v", a.BootstrapServerOnboardingInfo.IetfSztpConveyedInfoOnboardingInformation.BootImage.DownloadURI)
-	_ = a.doReportProgress(ProgressTypeBootImageInitiated, "BootImage Initiated")
-	// Download the image from DownloadURI and save it to a file
-	a.BootstrapServerOnboardingInfo.IetfSztpConveyedInfoOnboardingInformation.InfoTimestampReference = fmt.Sprintf("%8d", time.Now().Unix())
-	for i, item := range a.BootstrapServerOnboardingInfo.IetfSztpConveyedInfoOnboardingInformation.BootImage.DownloadURI {
-		// TODO: maybe need to file download to a function in util.go
-		log.Printf("[INFO] Downloading Image %v", item)
-		// Create a empty file
-		file, err := os.Create(ARTIFACTS_PATH + a.BootstrapServerOnboardingInfo.IetfSztpConveyedInfoOnboardingInformation.InfoTimestampReference + filepath.Base(item))
-		if err != nil {
-			return err
-		}
+func (a *Agent) downloadArtifact(uri string) (*os.File, error) {
+	file, err := os.CreateTemp("", filepath.Base(uri))
+	if err != nil {
+		return nil, err
+	}
 
-		caCert, _ := os.ReadFile(a.GetBootstrapTrustAnchorCert())
-		caCertPool := x509.NewCertPool()
-		caCertPool.AppendCertsFromPEM(caCert)
-		cert, _ := tls.LoadX509KeyPair(a.GetDeviceEndEntityCert(), a.GetDevicePrivateKey())
+	response, err := a.HttpClient.Get(uri)
+	if err != nil {
+		return nil, err
+	}
+	sizeorigin, _ := strconv.Atoi(response.Header.Get("Content-Length"))
+	downloadSize := int64(sizeorigin)
+	log.Printf("[INFO] Downloading the image with size: %v", downloadSize)
 
-		check := http.Client{
-			CheckRedirect: func(r *http.Request, _ []*http.Request) error {
-				r.URL.Opaque = r.URL.Path
-				return nil
-			},
-			Transport: &http.Transport{
-				TLSClientConfig: &tls.Config{
-					//nolint:gosec
-					InsecureSkipVerify: true, // TODO: remove skip verify
-					RootCAs:            caCertPool,
-					Certificates:       []tls.Certificate{cert},
-				},
-			},
+	if response.StatusCode != 200 {
+		return nil, errors.New("received non 200 response code")
+	}
+	size, err := io.Copy(file, response.Body)
+	if err != nil {
+		return nil, err
+	}
+	defer func() {
+		if err := file.Close(); err != nil {
+			log.Println("[ERROR] Error when closing:", err)
 		}
+	}()
+	defer func() {
+		if err := response.Body.Close(); err != nil {
+			log.Println("[ERROR] Error when closing:", err)
+		}
+	}()
+	log.Printf("[INFO] Downloaded file: %s with size: %d", file.Name(), size)
+	return file, nil
+}
+
+func (a *Agent) validateImage(filePath string, algorithm string, expected string) error {
+	switch algorithm {
+	case "ietf-sztp-conveyed-info:sha-256":
+		checksum, err := calculateSHA256File(filePath)
 
-		response, err := check.Get(item)
 		if err != nil {
-			return err
+			log.Println("[ERROR] Could not calculate checksum", err)
 		}
+		log.Println("calculated: " + checksum)
+		log.Println("expected  : " + expected)
 
-		sizeorigin, _ := strconv.Atoi(response.Header.Get("Content-Length"))
-		downloadSize := int64(sizeorigin)
-		log.Printf("[INFO] Downloading the image with size: %v", downloadSize)
-
-		if response.StatusCode != 200 {
-			return errors.New("received non 200 response code")
+		if checksum != expected {
+			return errors.New("checksum mismatch")
 		}
-		size, err := io.Copy(file, response.Body)
-		if err != nil {
-			return err
+		log.Println("[INFO] Checksum verified successfully")
+		return nil
+	default:
+		return errors.New("unsupported hash algorithm")
+	}
+}
+
+func (a *Agent) artifactExists(item string, algorithm string, expected string) bool {
+	filePath := ARTIFACTS_PATH + filepath.Base(item)
+	_, err := os.Stat(filePath)
+	if err != nil {
+		return false
+	}
+	err = a.validateImage(filePath, algorithm, expected)
+	return err == nil
+}
+
+//nolint:funlen
+func (a *Agent) downloadAndValidateImage() error {
+	bootImage := a.BootstrapServerOnboardingInfo.IetfSztpConveyedInfoOnboardingInformation.BootImage
+	log.Printf("[INFO] Starting the Download Image: %v", bootImage.DownloadURI)
+	_ = a.doReportProgress(ProgressTypeBootImageInitiated, "BootImage Initiated")
+	// Download the image from DownloadURI and save it to a file
+	for i, item := range bootImage.DownloadURI {
+		if len(bootImage.ImageVerification) <= i {
+			return errors.New("invalid verification")
 		}
-		defer func() {
-			if err := file.Close(); err != nil {
-				log.Println("[ERROR] Error when closing:", err)
-			}
-		}()
-		defer func() {
-			if err := response.Body.Close(); err != nil {
-				log.Println("[ERROR] Error when closing:", err)
-			}
-		}()
 
-		log.Printf("[INFO] Downloaded file: %s with size: %d", ARTIFACTS_PATH+a.BootstrapServerOnboardingInfo.IetfSztpConveyedInfoOnboardingInformation.InfoTimestampReference+filepath.Base(item), size)
-		log.Println("[INFO] Verify the file checksum: ", ARTIFACTS_PATH+a.BootstrapServerOnboardingInfo.IetfSztpConveyedInfoOnboardingInformation.InfoTimestampReference+filepath.Base(item))
-		switch a.BootstrapServerOnboardingInfo.IetfSztpConveyedInfoOnboardingInformation.BootImage.ImageVerification[i].HashAlgorithm {
-		case "ietf-sztp-conveyed-info:sha-256":
-			filePath := ARTIFACTS_PATH + a.BootstrapServerOnboardingInfo.IetfSztpConveyedInfoOnboardingInformation.InfoTimestampReference + filepath.Base(item)
-			checksum, err := calculateSHA256File(filePath)
-			original := strings.ReplaceAll(a.BootstrapServerOnboardingInfo.IetfSztpConveyedInfoOnboardingInformation.BootImage.ImageVerification[i].HashValue, ":", "")
+		imageVerification := bootImage.ImageVerification[i]
+		expected := strings.ReplaceAll(imageVerification.HashValue, ":", "")
+		algorithm := imageVerification.HashAlgorithm
+
+		if a.artifactExists(item, algorithm, expected) {
+			log.Printf("[INFO] Image %v already exists", item)
+		} else {
+			log.Printf("[INFO] Downloading Image %v", item)
+			file, err := a.downloadArtifact(item)
 			if err != nil {
-				log.Println("[ERROR] Could not calculate checksum", err)
+				return err
 			}
-			log.Println("calculated: " + checksum)
-			log.Println("expected  : " + original)
-			if checksum != original {
-				return errors.New("checksum mismatch")
+			log.Println("[INFO] Verify the file checksum: ", file.Name())
+			err = a.validateImage(file.Name(), algorithm, expected)
+			if err != nil {
+				return err
 			}
-			log.Println("[INFO] Checksum verified successfully")
+
+			log.Printf("[INFO] Moving file %s to %s", file.Name(), ARTIFACTS_PATH+filepath.Base(item))
+			_ = os.Rename(file.Name(), ARTIFACTS_PATH+filepath.Base(item))
 			_ = a.doReportProgress(ProgressTypeBootImageComplete, "BootImage Complete")
+
 			return nil
-		default:
-			return errors.New("unsupported hash algorithm")
 		}
 	}
 	return nil