COCOS-458 - Extend the CLI with the option to calculate the PCR16 register (#489)

* added code for extending policy json file

* minor bug fix

* typo

* added tests

* increased coverage

* fixed error

* improved coverage

* fixed according to comments

cli/attestation_policy.go

@@ -6,10 +6,12 @@ import (
 	"bytes"
 	"crypto/sha512"
 	"encoding/base64"
+	"encoding/hex"
 	"encoding/json"
 	"fmt"
 	"os"
 	"strconv"
+	"strings"
 
 	"github.com/absmach/supermq/pkg/errors"
 	"github.com/google/go-sev-guest/proto/check"
@@ -46,6 +48,8 @@ var (
 	errMarshalJSON                         = errors.New("failed to marshal json")
 	errWriteFile                           = errors.New("failed to write to file")
 	errAttestationPolicyField              = errors.New("the specified field type does not exist in the attestation policy")
+	errReadingManifestFile                 = errors.New("error while reading manifest file")
+	errDecodeHex                           = errors.New("error decoding hex string")
 	policy                          uint64 = 196639
 )
 
@@ -275,6 +279,23 @@ func (cli *CLI) NewAzureAttestationPolicy() *cobra.Command {
 	return cmd
 }
 
+func (cli *CLI) NewExtendWithManifestCmd() *cobra.Command {
+	return &cobra.Command{
+		Use:     "extend",
+		Short:   "Extends PCR16 with computation manifests. The first parameter is path to attestation policy file. The rest of the parameters are paths to computation manifest files.",
+		Example: "extend <attestation_policy_file_path> <computation_manifest_file_path> [<computation_manifest_file_path> ...]",
+		Args:    cobra.MinimumNArgs(2),
+		Run: func(cmd *cobra.Command, args []string) {
+			attestationPolicyFilePath := args[0]
+			manifestPaths := args[1:]
+			if err := extendWithManifest(attestationPolicyFilePath, manifestPaths); err != nil {
+				printError(cmd, "Error could not extend PCR16: %v ❌ ", err)
+				return
+			}
+		},
+	}
+}
+
 func changeAttestationConfiguration(fileName, base64Data string, expectedLength int, field fieldType) error {
 	data, err := base64.StdEncoding.DecodeString(base64Data)
 	if err != nil {
@@ -318,3 +339,65 @@ func changeAttestationConfiguration(fileName, base64Data string, expectedLength
 	}
 	return nil
 }
+
+func extendWithManifest(attestationPolicyPath string, manifestPaths []string) error {
+	attestationConfig := attestation.Config{Config: &check.Config{RootOfTrust: &check.RootOfTrust{}, Policy: &check.Policy{}}, PcrConfig: &attestation.PcrConfig{}}
+
+	attestationPolicyFileData, err := os.ReadFile(attestationPolicyPath)
+	if err != nil {
+		return errors.Wrap(errReadingAttestationPolicyFile, err)
+	}
+
+	if err = vtpm.ReadPolicyFromByte(attestationPolicyFileData, &attestationConfig); err != nil {
+		return errors.Wrap(errUnmarshalJSON, err)
+	}
+
+	for _, manifestPath := range manifestPaths {
+		manifest, err := os.ReadFile(manifestPath)
+		if err != nil {
+			return errors.Wrap(errReadingManifestFile, err)
+		}
+
+		manifestSha256 := sha512.Sum512_256(manifest)
+		manifestSha384 := sha512.Sum384(manifest)
+
+		data256, exists256 := attestationConfig.PCRValues.Sha256["16"]
+
+		if !exists256 {
+			data256 = strings.Repeat("0", 64) // 32 bytes in hex
+		}
+
+		byteData256, err := hex.DecodeString(data256)
+		if err != nil {
+			return errors.Wrap(errDecodeHex, err)
+		}
+
+		newByteData256 := sha512.Sum512_256(append(byteData256, manifestSha256[:]...))
+
+		data384, exists384 := attestationConfig.PCRValues.Sha384["16"]
+
+		if !exists384 {
+			data384 = strings.Repeat("0", 96) // 48 bytes in hex
+		}
+
+		byteData384, err := hex.DecodeString(data384)
+		if err != nil {
+			return errors.Wrap(errDecodeHex, err)
+		}
+
+		newByteData384 := sha512.Sum384(append(byteData384, manifestSha384[:]...))
+
+		attestationConfig.PCRValues.Sha256["16"] = hex.EncodeToString(newByteData256[:])
+		attestationConfig.PCRValues.Sha384["16"] = hex.EncodeToString(newByteData384[:])
+	}
+
+	attestationPolicyJSON, err := vtpm.ConvertPolicyToJSON(&attestationConfig)
+	if err != nil {
+		return errors.Wrap(errMarshalJSON, err)
+	}
+	if err = os.WriteFile(attestationPolicyPath, attestationPolicyJSON, filePermission); err != nil {
+		return errors.Wrap(errWriteFile, err)
+	}
+
+	return nil
+}

cli/attestation_policy_test.go

@@ -376,3 +376,96 @@ func TestCommandErrorHandling(t *testing.T) {
 		assert.Contains(t, output, "❌")
 	})
 }
+
+func TestExtendWithManifestHandling(t *testing.T) {
+	cli := &CLI{}
+
+	t.Run("Invalid policy file", func(t *testing.T) {
+		cmd := cli.NewExtendWithManifestCmd()
+		cmd.SetArgs([]string{"nonexistent.policy.json", "nonexistent.manifest.json"})
+
+		var buf bytes.Buffer
+		cmd.SetOut(&buf)
+		cmd.SetErr(&buf)
+
+		err := cmd.Execute()
+		assert.NoError(t, err)
+
+		output := buf.String()
+		assert.Contains(t, output, "error while reading the attestation policy file")
+		assert.Contains(t, output, "❌")
+	})
+
+	t.Run("Invalid manifest file", func(t *testing.T) {
+		cmd := cli.NewExtendWithManifestCmd()
+		cmd.SetArgs([]string{"../scripts/attestation_policy/attestation_policy.json", "nonexistent.manifest.json"})
+
+		var buf bytes.Buffer
+		cmd.SetOut(&buf)
+		cmd.SetErr(&buf)
+
+		err := cmd.Execute()
+		assert.NoError(t, err)
+
+		output := buf.String()
+		assert.Contains(t, output, "error while reading manifest file")
+		assert.Contains(t, output, "❌")
+	})
+
+	t.Run("Valid file paths", func(t *testing.T) {
+		fileContent := `{
+		  "id": "1",
+		  "name": "sample computation",
+		  "description": "sample description",
+		  "datasets": [
+		    {
+		      "hash": "<sha3_encoded string>",
+		      "userKey": "<pem_encoded public key string>"
+		    }
+		  ],
+		  "algorithm": {
+		    "hash": "<sha3_encoded string>",
+		    "userKey": "<pem_encoded public key string>"
+		  },
+		  "result_consumers": [
+		    {
+		      "userKey": "<pem_encoded public key string>"
+		    }
+		  ],
+		  "agent_config": {
+		    "port": "7002",
+		    "cert_file": "<pem encoded cert string>",
+		    "key_file": "<pem encoded private key string>",
+		    "server_ca_file": "<pem encoded cert string>",
+		    "client_ca_file": "<pem encoded cert string>",
+		    "attested_tls": true
+		  }
+		}`
+
+		dir, err := os.Getwd()
+		if err != nil {
+			t.Fatalf("Error getting current working directory: %v", err)
+		}
+
+		manifestFile, err := os.CreateTemp(dir, "manifest.json")
+		if err != nil {
+			t.Fatalf("Error creating temp file: %v", err)
+		}
+		defer os.Remove(manifestFile.Name())
+
+		err = os.WriteFile(manifestFile.Name(), []byte(fileContent), 0o644)
+		if err != nil {
+			t.Fatalf("Error writing temp file: %v", err)
+		}
+
+		cmd := cli.NewExtendWithManifestCmd()
+		cmd.SetArgs([]string{"../scripts/attestation_policy/attestation_policy.json", manifestFile.Name()})
+
+		var buf bytes.Buffer
+		cmd.SetOut(&buf)
+		cmd.SetErr(&buf)
+
+		err = cmd.Execute()
+		assert.NoError(t, err)
+	})
+}