From 8376c579913fdadca280557af8d94280642e7ba5 Mon Sep 17 00:00:00 2001
From: Boris Glimcher <Boris.Glimcher@emc.com>
Date: Sat, 3 Aug 2024 00:33:37 +0300
Subject: [PATCH] refactor: move CalculateFileSHA256 to utils

Signed-off-by: Boris Glimcher <Boris.Glimcher@emc.com>
---
 sztp-agent/pkg/secureagent/daemon.go     | 14 +-------
 sztp-agent/pkg/secureagent/utils.go      | 28 +++++++++++++++
 sztp-agent/pkg/secureagent/utils_test.go | 44 ++++++++++++++++++++++++
 3 files changed, 73 insertions(+), 13 deletions(-)

diff --git a/sztp-agent/pkg/secureagent/daemon.go b/sztp-agent/pkg/secureagent/daemon.go
index e646592..523796b 100644
--- a/sztp-agent/pkg/secureagent/daemon.go
+++ b/sztp-agent/pkg/secureagent/daemon.go
@@ -10,7 +10,6 @@ package secureagent
 
 import (
 	"bytes"
-	"crypto/sha256"
 	"crypto/tls"
 	"crypto/x509"
 	"encoding/asn1"
@@ -249,21 +248,10 @@ func (a *Agent) downloadAndValidateImage() error {
 		// TODO: maybe need to move sha calculatinos to a function in util.go
 		switch a.BootstrapServerOnboardingInfo.IetfSztpConveyedInfoOnboardingInformation.BootImage.ImageVerification[i].HashAlgorithm {
 		case "ietf-sztp-conveyed-info:sha-256":
-			f, err := os.Open(ARTIFACTS_PATH + a.BootstrapServerOnboardingInfo.IetfSztpConveyedInfoOnboardingInformation.InfoTimestampReference + filepath.Base(item))
+			sum, err := CalculateFileSHA256(ARTIFACTS_PATH + a.BootstrapServerOnboardingInfo.IetfSztpConveyedInfoOnboardingInformation.InfoTimestampReference + filepath.Base(item))
 			if err != nil {
-				log.Panic(err)
 				return err
 			}
-			defer func() {
-				if err := f.Close(); err != nil {
-					log.Println("[ERROR] Error when closing:", err)
-				}
-			}()
-			h := sha256.New()
-			if _, err := io.Copy(h, f); err != nil {
-				return err
-			}
-			sum := fmt.Sprintf("%x", h.Sum(nil))
 			original := strings.ReplaceAll(a.BootstrapServerOnboardingInfo.IetfSztpConveyedInfoOnboardingInformation.BootImage.ImageVerification[i].HashValue, ":", "")
 			log.Println("calculated: " + sum)
 			log.Println("expected  : " + original)
diff --git a/sztp-agent/pkg/secureagent/utils.go b/sztp-agent/pkg/secureagent/utils.go
index b8cf22c..87772f6 100644
--- a/sztp-agent/pkg/secureagent/utils.go
+++ b/sztp-agent/pkg/secureagent/utils.go
@@ -9,8 +9,12 @@ Copyright (C) 2022 Red Hat.
 package secureagent
 
 import (
+	"crypto/sha256"
 	"encoding/json"
+	"fmt"
+	"io"
 	"log"
+	"os"
 	"strings"
 
 	"github.com/go-ini/ini"
@@ -34,6 +38,30 @@ func GetSerialNumber(givenSerialNumber string) string {
 	return serialNumber
 }
 
+// CalculateFileSHA256 computes the SHA-256 checksum of a file specified by its path.
+func CalculateFileSHA256(filePath string) (string, error) {
+	// Open the file
+	file, err := os.Open(filePath)
+	if err != nil {
+		return "", err
+	}
+	defer file.Close()
+
+	// Create a new SHA256 hash object
+	hash := sha256.New()
+
+	// Copy the file content to the hash object
+	if _, err := io.Copy(hash, file); err != nil {
+		return "", err
+	}
+
+	// Get the final SHA256 hash result
+	checksum := hash.Sum(nil)
+
+	// Convert the result to a hexadecimal string
+	return fmt.Sprintf("%x", checksum), nil
+}
+
 func generateInputJSONContent() string {
 	osName := ""
 	osVersion := ""
diff --git a/sztp-agent/pkg/secureagent/utils_test.go b/sztp-agent/pkg/secureagent/utils_test.go
index a3cd990..7da2070 100644
--- a/sztp-agent/pkg/secureagent/utils_test.go
+++ b/sztp-agent/pkg/secureagent/utils_test.go
@@ -5,6 +5,7 @@
 package secureagent
 
 import (
+	"os"
 	"testing"
 )
 
@@ -47,3 +48,46 @@ func Test_replaceQuotes(t *testing.T) {
 		})
 	}
 }
+
+func TestCalculateFileSHA256(t *testing.T) {
+	// Define test cases
+	tests := []struct {
+		name     string
+		filePath string
+		want     string
+		wantErr  bool
+	}{
+		{
+			name:     "Valid file",
+			filePath: "testfile.txt",
+			want:     "d2d2d2d2d2d2d2d2d2d2d2d2d2d2d2d2d2d2d2d2d2d2d2d2d2d2d2d2d2d2d2d2d2d2d2d2d2", // Replace with actual expected hash
+			wantErr:  false,
+		},
+		{
+			name:     "Non-existent file",
+			filePath: "nonexistentfile.txt",
+			want:     "",
+			wantErr:  true,
+		},
+	}
+
+	// Create a valid file for testing
+	err := os.WriteFile("testfile.txt", []byte("test"), 0644)
+	if err != nil {
+		t.Fatalf("Failed to create test file: %v", err)
+	}
+	defer os.Remove("testfile.txt") // Clean up the file after tests
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := CalculateFileSHA256(tt.filePath)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("CalculateFileSHA256() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if got != tt.want {
+				t.Errorf("CalculateFileSHA256() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}