COCOS-253 - Improve CLI error handling (#277)

* decode errors Signed-off-by: Sammy Oina <sammyoina@gmail.com> * standardise error formatting Signed-off-by: Sammy Oina <sammyoina@gmail.com> * fix failing tests Signed-off-by: Sammy Oina <sammyoina@gmail.com> * add errors tests Signed-off-by: Sammy Oina <sammyoina@gmail.com> * pass lint Signed-off-by: Sammy Oina <sammyoina@gmail.com> * add test cases Signed-off-by: Sammy Oina <sammyoina@gmail.com> --------- Signed-off-by: Sammy Oina <sammyoina@gmail.com>
2026-06-23 04:10:25 +00:00 · 2024-10-08 18:11:37 +03:00
parent 7ef25674c4
commit fb0fbaeb9a
17 changed files with 491 additions and 256 deletions
@@ -9,7 +9,6 @@ import (
 	"crypto/x509"
 	"encoding/pem"
 	"errors"
-	"log"
 	"os"
 	"testing"

@@ -20,14 +19,6 @@ import (

 const algorithmFile = "test_algo_file.py"

-func captureLogOutput(f func()) string {
-	var buf bytes.Buffer
-	log.SetOutput(&buf)
-	defer log.SetOutput(os.Stderr)
-	f()
-	return buf.String()
-}
-
 func generateRSAPrivateKeyFile(fileName string) error {
 	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
 	if err != nil {
@@ -65,13 +56,13 @@ func TestAlgorithmCmd_Success(t *testing.T) {
 	require.NoError(t, err)

 	cmd := testCLI.NewAlgorithmCmd()
-	output := captureLogOutput(func() {
-		cmd.SetArgs([]string{algorithmFile, privateKeyFile})
-		err = cmd.Execute()
-		require.NoError(t, err)
-	})
+	buf := new(bytes.Buffer)
+	cmd.SetOut(buf)
+	cmd.SetArgs([]string{algorithmFile, privateKeyFile})
+	err = cmd.Execute()
+	require.NoError(t, err)

-	require.Contains(t, output, "Successfully uploaded algorithm")
+	require.Contains(t, buf.String(), "Successfully uploaded algorithm")
 	t.Cleanup(func() {
 		os.Remove(privateKeyFile)
 		os.Remove(algorithmFile)
@@ -84,14 +75,14 @@ func TestAlgorithmCmd_MissingAlgorithmFile(t *testing.T) {
 	testCLI := New(mockSDK)

 	cmd := testCLI.NewAlgorithmCmd()
+	buf := new(bytes.Buffer)
+	cmd.SetOut(buf)

-	output := captureLogOutput(func() {
-		cmd.SetArgs([]string{"non_existent_algo_file.py", privateKeyFile})
-		err := cmd.Execute()
-		require.NoError(t, err)
-	})
+	cmd.SetArgs([]string{"non_existent_algo_file.py", privateKeyFile})
+	err := cmd.Execute()
+	require.NoError(t, err)

-	require.Contains(t, output, "Error reading algorithm file")
+	require.Contains(t, buf.String(), "Error reading algorithm file")
 }

 func TestAlgorithmCmd_MissingPrivateKeyFile(t *testing.T) {
@@ -103,14 +94,13 @@ func TestAlgorithmCmd_MissingPrivateKeyFile(t *testing.T) {
 	require.NoError(t, err)

 	cmd := testCLI.NewAlgorithmCmd()
+	buf := new(bytes.Buffer)
+	cmd.SetOut(buf)
+	cmd.SetArgs([]string{algorithmFile, "non_existent_private_key.pem"})
+	err = cmd.Execute()
+	require.NoError(t, err)

-	output := captureLogOutput(func() {
-		cmd.SetArgs([]string{algorithmFile, "non_existent_private_key.pem"})
-		err = cmd.Execute()
-		require.NoError(t, err)
-	})
-
-	require.Contains(t, output, "Error reading private key file")
+	require.Contains(t, buf.String(), "Error reading private key file")
 	t.Cleanup(func() {
 		os.Remove(algorithmFile)
 	})
@@ -128,17 +118,48 @@ func TestAlgorithmCmd_UploadFailure(t *testing.T) {
 	require.NoError(t, err)

 	cmd := testCLI.NewAlgorithmCmd()
+	buf := new(bytes.Buffer)
+	cmd.SetOut(buf)

-	output := captureLogOutput(func() {
-		cmd.SetArgs([]string{algorithmFile, privateKeyFile})
-		err = cmd.Execute()
-		require.NoError(t, err)
-	})
+	cmd.SetArgs([]string{algorithmFile, privateKeyFile})
+	err = cmd.Execute()
+	require.NoError(t, err)

-	require.Contains(t, output, "Failed to upload algorithm")
+	require.Contains(t, buf.String(), "Failed to upload algorithm")

 	t.Cleanup(func() {
 		os.Remove(privateKeyFile)
 		os.Remove(algorithmFile)
 	})
 }
+
+func TestAlgorithmCmd_InvalidPrivateKey(t *testing.T) {
+	mockSDK := new(mocks.SDK)
+	mockSDK.On("Algo", mock.Anything, mock.Anything, mock.Anything).Return(nil)
+	testCLI := New(mockSDK)
+
+	err := os.WriteFile(algorithmFile, []byte("test algorithm"), 0o644)
+	require.NoError(t, err)
+
+	privKeyFile, err := os.Create(privateKeyFile)
+	require.NoError(t, err)
+	defer privKeyFile.Close()
+
+	_, err = privKeyFile.WriteString("invalid private key")
+	require.NoError(t, err)
+
+	cmd := testCLI.NewAlgorithmCmd()
+	buf := new(bytes.Buffer)
+	cmd.SetOut(buf)
+
+	cmd.SetArgs([]string{algorithmFile, privateKeyFile})
+	err = cmd.Execute()
+	require.NoError(t, err)
+
+	require.Contains(t, buf.String(), "Error decoding private key")
+
+	t.Cleanup(func() {
+		os.Remove(algorithmFile)
+		os.Remove(privateKeyFile)
+	})
+}
@@ -5,7 +5,6 @@ package cli
 import (
 	"context"
 	"encoding/pem"
-	"log"
 	"os"

 	"github.com/fatih/color"
@@ -32,12 +31,11 @@ func (cli *CLI) NewAlgorithmCmd() *cobra.Command {
 		Run: func(cmd *cobra.Command, args []string) {
 			algorithmFile := args[0]

-			log.Println("Uploading algorithm file:", algorithmFile)
+			cmd.Println("Uploading algorithm file:", algorithmFile)

 			algorithm, err := os.ReadFile(algorithmFile)
 			if err != nil {
-				msg := color.New(color.FgRed).Sprintf("Error reading algorithm file: %v ❌ ", err)
-				log.Println(msg)
+				printError(cmd, "Error reading algorithm file: %v ❌ ", err)
 				return
 			}

@@ -45,8 +43,7 @@ func (cli *CLI) NewAlgorithmCmd() *cobra.Command {
 			if requirementsFile != "" {
 				req, err = os.ReadFile(requirementsFile)
 				if err != nil {
-					msg := color.New(color.FgRed).Sprintf("Error reading requirments file: %v ❌ ", err)
-					log.Println(msg)
+					printError(cmd, "Error reading requirments file: %v ❌ ", err)
 					return
 				}
 			}
@@ -58,24 +55,26 @@ func (cli *CLI) NewAlgorithmCmd() *cobra.Command {

 			privKeyFile, err := os.ReadFile(args[1])
 			if err != nil {
-				msg := color.New(color.FgRed).Sprintf("Error reading private key file: %v ❌ ", err.Error())
-				log.Println(msg)
+				printError(cmd, "Error reading private key file: %v ❌ ", err)
 				return
 			}

 			pemBlock, _ := pem.Decode(privKeyFile)

-			privKey := decodeKey(pemBlock)
+			privKey, err := decodeKey(pemBlock)
+			if err != nil {
+				printError(cmd, "Error decoding private key: %v ❌ ", err)
+				return
+			}

 			ctx := metadata.NewOutgoingContext(cmd.Context(), metadata.New(make(map[string]string)))

 			if err := cli.agentSDK.Algo(addAlgoMetadata(ctx), algoReq, privKey); err != nil {
-				msg := color.New(color.FgRed).Sprintf("Failed to upload algorithm due to error: %v ❌ ", err.Error())
-				log.Println(msg)
+				printError(cmd, "Failed to upload algorithm due to error: %v ❌ ", err)
 				return
 			}

-			log.Println(color.New(color.FgGreen).Sprint("Successfully uploaded algorithm! ✔ "))
+			cmd.Println(color.New(color.FgGreen).Sprint("Successfully uploaded algorithm! ✔ "))
 		},
 	}

@@ -5,12 +5,12 @@ package cli
 import (
 	"encoding/hex"
 	"fmt"
-	"log"
 	"os"
 	"strconv"
 	"strings"
 	"time"

+	"github.com/fatih/color"
 	"github.com/google/go-sev-guest/abi"
 	"github.com/google/go-sev-guest/proto/check"
 	"github.com/google/go-sev-guest/proto/sevsnp"
@@ -158,22 +158,23 @@ func (cli *CLI) NewGetAttestationCmd() *cobra.Command {

 			reportData, err := hex.DecodeString(args[0])
 			if err != nil {
-				cmd.Printf("attestation validation and verification failed with error: %s", err)
+				printError(cmd, "Error decoding report data: %v ❌ ", err)
 				return
 			}
 			if len(reportData) != agent.ReportDataSize {
-				cmd.Printf("report data must be a hex encoded string of length %d bytes", agent.ReportDataSize)
+				msg := color.New(color.FgRed).Sprintf("report data must be a hex encoded string of length %d bytes ❌ ", agent.ReportDataSize)
+				cmd.Println(msg)
 				return
 			}

 			result, err := cli.agentSDK.Attestation(cmd.Context(), [agent.ReportDataSize]byte(reportData))
 			if err != nil {
-				cmd.Printf("Error retrieving attestation: %v", err)
+				printError(cmd, "Failed to get attestation due to error: %v ❌ ", err)
 				return
 			}

 			if err = os.WriteFile(attestationFilePath, result, 0o644); err != nil {
-				cmd.Printf("Error saving attestation result: %v", err)
+				printError(cmd, "Error saving attestation result: %v ❌ ", err)
 				return
 			}

@@ -189,37 +190,45 @@ func (cli *CLI) NewValidateAttestationValidationCmd() *cobra.Command {
 		Example: "validate <attestation_report_file_path>",
 		Args:    cobra.ExactArgs(1),
 		Run: func(cmd *cobra.Command, args []string) {
-			log.Println("Checking attestation")
+			cmd.Println("Checking attestation")

 			attestationFile = string(args[0])

 			if err := parseConfig(); err != nil {
-				log.Fatalf("attestation validation and verification failed with error: %s", err)
+				printError(cmd, "Error parsing config: %v ❌ ", err)
+				return
 			}
 			if err := parseHashes(); err != nil {
-				log.Fatalf("attestation validation and verification failed with error: %s", err)
+				printError(cmd, "Error parsing hashes: %v ❌ ", err)
+				return
 			}
 			if err := parseFiles(); err != nil {
-				log.Fatalf("attestation validation and verification failed with error: %s", err)
+				printError(cmd, "Error parsing files: %v ❌ ", err)
+				return
 			}
 			// This format is the attestation report in AMD's specified ABI format, immediately
 			// followed by the certificate table bytes.
 			if len(attestation) < abi.ReportSize {
-				log.Fatalf("attestation contents too small (0x%x bytes). Want at least 0x%x bytes", len(attestation), abi.ReportSize)
+				msg := color.New(color.FgRed).Sprintf("attestation contents too small (0x%x bytes). Want at least 0x%x bytes ❌ ", len(attestation), abi.ReportSize)
+				cmd.Println(msg)
+				return
 			}
 			if err := parseUints(); err != nil {
-				log.Fatalf("attestation validation and verification failed with error: %s", err)
+				printError(cmd, "Error parsing uints: %v ❌ ", err)
+				return
 			}
 			cfg.Policy.Vmpl = wrapperspb.UInt32(0)

 			if err := validateInput(); err != nil {
-				log.Fatalf("attestation validation and verification failed with error: %s", err)
+				printError(cmd, "Error validating input: %v ❌ ", err)
+				return
 			}

 			if err := verifyAndValidateAttestation(attestation); err != nil {
-				log.Fatalf("attestation validation and verification failed with error: %s", err)
+				printError(cmd, "Attestation validation and verification failed with error: %v ❌ ", err)
+				return
 			}
-			log.Println("Attestation validation and verification is successful!")
+			cmd.Println("Attestation validation and verification is successful!")
 		},
 	}
 	cmd.Flags().StringVar(
@@ -398,11 +407,13 @@ func (cli *CLI) NewValidateAttestationValidationCmd() *cobra.Command {
 	)

 	if err := cmd.MarkFlagRequired("report_data"); err != nil {
-		log.Fatalf("Failed to mark flag as required: %s", err)
+		printError(cmd, "Failed to mark flag as required: %v ❌ ", err)
+		return nil
 	}

 	if err := cmd.MarkFlagRequired("product"); err != nil {
-		log.Fatalf("Failed to mark flag as required: %s", err)
+		printError(cmd, "Failed to mark flag as required: %v ❌ ", err)
+		return nil
 	}

 	return cmd
@@ -6,7 +6,6 @@ import (
 	"encoding/base64"
 	"encoding/json"
 	"fmt"
-	"log"
 	"os"

 	"github.com/absmach/magistrala/pkg/errors"
@@ -83,7 +82,8 @@ func (cli *CLI) NewAddMeasurementCmd() *cobra.Command {
 		Args:    cobra.ExactArgs(2),
 		Run: func(cmd *cobra.Command, args []string) {
 			if err := changeAttestationConfiguration(args[1], args[0], measurementLength, measurementField); err != nil {
-				log.Fatalf("Error could not change measurement data %v", err)
+				printError(cmd, "Error could not change measurement data: %v ❌ ", err)
+				return
 			}
 		},
 	}
@@ -97,7 +97,8 @@ func (cli *CLI) NewAddHostDataCmd() *cobra.Command {
 		Args:    cobra.ExactArgs(2),
 		Run: func(cmd *cobra.Command, args []string) {
 			if err := changeAttestationConfiguration(args[1], args[0], hostDataLength, hostDataField); err != nil {
-				log.Fatalf("Error could not change host data %v", err)
+				printError(cmd, "Error could not change host data: %v ❌ ", err)
+				return
 			}
 		},
 	}
@@ -3,7 +3,7 @@
 package cli

 import (
-	"log"
+	"fmt"
 	"os"
 	"path"

@@ -29,7 +29,8 @@ func (cli *CLI) NewCABundleCmd(fileSavePath string) *cobra.Command {
 			attestationConfiguration := grpc.AttestationConfiguration{}
 			err := grpc.ReadBackendInfo(args[0], &attestationConfiguration)
 			if err != nil {
-				log.Fatalf("Error while reading manifest: %v", err)
+				printError(cmd, "Error while reading manifest: %v ❌ ", err)
+				return
 			}

 			product := attestationConfiguration.RootOfTrust.Product
@@ -39,17 +40,24 @@ func (cli *CLI) NewCABundleCmd(fileSavePath string) *cobra.Command {

 			bundle, err := getter.Get(caURL)
 			if err != nil {
-				log.Fatalf("Error fetching ARK and ASK from AMD KDS for product: %s, error: %v", product, err)
+				message := fmt.Sprintf("Error fetching ARK and ASK from AMD KDS for product: %s", product)
+				message += ", error: %v ❌ "
+				printError(cmd, message, err)
+				return
 			}

 			err = os.MkdirAll(path.Join(fileSavePath, product), filePermisionKeys)
 			if err != nil {
-				log.Fatalf("Error while creating directory for product name %s, error: %v", product, err)
+				message := fmt.Sprintf("Error while creating directory for product name %s", product)
+				message += ", error: %v ❌ "
+				printError(cmd, message, err)
+				return
 			}

 			bundleFilePath := path.Join(fileSavePath, product, caBundleName)
 			if err = saveToFile(bundleFilePath, bundle); err != nil {
-				log.Fatalf("Error while saving ARK-ASK to file: %v", err)
+				printError(cmd, "Error while saving ARK-ASK to file: %v ❌ ", err)
+				return
 			}
 		},
 	}
@@ -18,7 +18,7 @@ func (cli *CLI) NewFileHashCmd() *cobra.Command {

 			hash, err := internal.ChecksumHex(path)
 			if err != nil {
-				cmd.Printf("Error computing hash: %v", err)
+				printError(cmd, "Error computing hash: %v ❌ ", err)
 				return
 			}

@@ -6,10 +6,10 @@ import (
 	"context"
 	"crypto/x509"
 	"encoding/pem"
-	"log"
 	"os"
 	"path"

+	"github.com/absmach/magistrala/pkg/errors"
 	"github.com/fatih/color"
 	"github.com/spf13/cobra"
 	"github.com/ultravioletrs/cocos/agent"
@@ -28,12 +28,11 @@ func (cli *CLI) NewDatasetsCmd() *cobra.Command {
 		Run: func(cmd *cobra.Command, args []string) {
 			datasetPath := args[0]

-			log.Println("Uploading dataset:", datasetPath)
+			cmd.Println("Uploading dataset:", datasetPath)

 			f, err := os.Stat(datasetPath)
 			if err != nil {
-				msg := color.New(color.FgRed).Sprintf("Error reading dataset file: %v ❌ ", err)
-				log.Println(msg)
+				printError(cmd, "Error reading dataset file: %v ❌ ", err)
 				return
 			}

@@ -42,15 +41,13 @@ func (cli *CLI) NewDatasetsCmd() *cobra.Command {
 			if f.IsDir() {
 				dataset, err = internal.ZipDirectoryToMemory(datasetPath)
 				if err != nil {
-					msg := color.New(color.FgRed).Sprintf("Error zipping dataset directory: %v ❌ ", err)
-					log.Println(msg)
+					printError(cmd, "Error zipping dataset directory: %v ❌ ", err)
 					return
 				}
 			} else {
 				dataset, err = os.ReadFile(datasetPath)
 				if err != nil {
-					msg := color.New(color.FgRed).Sprintf("Error reading dataset file: %v ❌ ", err)
-					log.Println(msg)
+					printError(cmd, "Error reading dataset file: %v ❌ ", err)
 					return
 				}
 			}
@@ -62,23 +59,25 @@ func (cli *CLI) NewDatasetsCmd() *cobra.Command {

 			privKeyFile, err := os.ReadFile(args[1])
 			if err != nil {
-				msg := color.New(color.FgRed).Sprintf("Error reading private key file: %v ❌ ", err)
-				log.Println(msg)
+				printError(cmd, "Error reading private key file: %v ❌ ", err)
 				return
 			}

 			pemBlock, _ := pem.Decode(privKeyFile)

-			privKey := decodeKey(pemBlock)
-
-			ctx := metadata.NewOutgoingContext(cmd.Context(), metadata.New(make(map[string]string)))
-			if err := cli.agentSDK.Data(addDatasetMetadata(ctx), dataReq, privKey); err != nil {
-				msg := color.New(color.FgRed).Sprintf("Failed to upload dataset due to error: %v ❌ ", err.Error())
-				log.Println(msg)
+			privKey, err := decodeKey(pemBlock)
+			if err != nil {
+				printError(cmd, "Error decoding private key: %v ❌ ", err)
 				return
 			}

-			log.Println(color.New(color.FgGreen).Sprint("Successfully uploaded dataset! ✔ "))
+			ctx := metadata.NewOutgoingContext(cmd.Context(), metadata.New(make(map[string]string)))
+			if err := cli.agentSDK.Data(addDatasetMetadata(ctx), dataReq, privKey); err != nil {
+				printError(cmd, "Failed to upload dataset due to error: %v ❌ ", err)
+				return
+			}
+
+			cmd.Println(color.New(color.FgGreen).Sprint("Successfully uploaded dataset! ✔ "))
 		},
 	}

@@ -86,26 +85,28 @@ func (cli *CLI) NewDatasetsCmd() *cobra.Command {
 	return cmd
 }

-func decodeKey(b *pem.Block) interface{} {
+func decodeKey(b *pem.Block) (interface{}, error) {
+	if b == nil {
+		return nil, errors.New("error decoding key")
+	}
 	switch b.Type {
 	case rsaKeyType:
 		privKey, err := x509.ParsePKCS8PrivateKey(b.Bytes)
 		if err != nil {
 			privKey, err = x509.ParsePKCS1PrivateKey(b.Bytes)
 			if err != nil {
-				log.Fatalf("Error parsing private key: %v", err)
+				return nil, err
 			}
 		}
-		return privKey
+		return privKey, nil
 	case ecdsaKeyType:
 		privKey, err := x509.ParseECPrivateKey(b.Bytes)
 		if err != nil {
-			log.Fatalf("Error parsing private key: %v", err)
+			return nil, err
 		}
-		return privKey
+		return privKey, nil
 	default:
-		log.Fatalf("Error decoding key")
-		return nil
+		return nil, errors.New("error decoding key")
 	}
 }

@@ -3,6 +3,7 @@
 package cli

 import (
+	"bytes"
 	"errors"
 	"os"
 	"testing"
@@ -38,14 +39,14 @@ func TestDatasetsCmd_Success(t *testing.T) {
 	require.NoError(t, err)

 	cmd := testCLI.NewDatasetsCmd()
+	buf := new(bytes.Buffer)
+	cmd.SetOut(buf)

-	output := captureLogOutput(func() {
-		cmd.SetArgs([]string{datasetFile, privateKeyFile})
-		err = cmd.Execute()
-		require.NoError(t, err)
-	})
+	cmd.SetArgs([]string{datasetFile, privateKeyFile})
+	err = cmd.Execute()
+	require.NoError(t, err)

-	require.Contains(t, output, "Successfully uploaded dataset")
+	require.Contains(t, buf.String(), "Successfully uploaded dataset")
 	mockSDK.AssertCalled(t, "Data", mock.Anything, mock.Anything, mock.Anything)

 	t.Cleanup(func() {
@@ -60,14 +61,14 @@ func TestDatasetsCmd_MissingDatasetFile(t *testing.T) {
 	testCLI := New(mockSDK)

 	cmd := testCLI.NewDatasetsCmd()
+	buf := new(bytes.Buffer)
+	cmd.SetOut(buf)

-	output := captureLogOutput(func() {
-		cmd.SetArgs([]string{"non_existent_dataset.txt", privateKeyFile})
-		err := cmd.Execute()
-		require.NoError(t, err)
-	})
+	cmd.SetArgs([]string{"non_existent_dataset.txt", privateKeyFile})
+	err := cmd.Execute()
+	require.NoError(t, err)

-	require.Contains(t, output, "Error reading dataset file")
+	require.Contains(t, buf.String(), "Error reading dataset file")
 }

 func TestDatasetsCmd_MissingPrivateKeyFile(t *testing.T) {
@@ -79,14 +80,14 @@ func TestDatasetsCmd_MissingPrivateKeyFile(t *testing.T) {
 	require.NoError(t, err)

 	cmd := testCLI.NewDatasetsCmd()
+	buf := new(bytes.Buffer)
+	cmd.SetOut(buf)

-	output := captureLogOutput(func() {
-		cmd.SetArgs([]string{datasetFile, "non_existent_private_key.pem"})
-		err = cmd.Execute()
-		require.NoError(t, err)
-	})
+	cmd.SetArgs([]string{datasetFile, "non_existent_private_key.pem"})
+	err = cmd.Execute()
+	require.NoError(t, err)

-	require.Contains(t, output, "Error reading private key file")
+	require.Contains(t, buf.String(), "Error reading private key file")
 	t.Cleanup(func() {
 		os.Remove(datasetFile)
 	})
@@ -104,14 +105,40 @@ func TestDatasetsCmd_UploadFailure(t *testing.T) {
 	require.NoError(t, err)

 	cmd := testCLI.NewDatasetsCmd()
+	buf := new(bytes.Buffer)
+	cmd.SetOut(buf)

-	output := captureLogOutput(func() {
-		cmd.SetArgs([]string{datasetFile, privateKeyFile})
-		err = cmd.Execute()
-		require.NoError(t, err)
-	})
+	cmd.SetArgs([]string{datasetFile, privateKeyFile})
+	err = cmd.Execute()
+	require.NoError(t, err)

-	require.Contains(t, output, "Failed to upload dataset due to error")
+	require.Contains(t, buf.String(), "Failed to upload dataset due to error")
+	t.Cleanup(func() {
+		os.Remove(datasetFile)
+		os.Remove(privateKeyFile)
+	})
+}
+
+func TestDatasetsCmd_InvalidPrivateKey(t *testing.T) {
+	mockSDK := new(mocks.SDK)
+	mockSDK.On("Data", mock.Anything, mock.Anything, mock.Anything).Return(nil)
+	testCLI := New(mockSDK)
+
+	datasetFile, err := createTempDatasetFile("test dataset content")
+	require.NoError(t, err)
+
+	err = os.WriteFile(privateKeyFile, []byte("invalid private key"), 0o644)
+	require.NoError(t, err)
+
+	cmd := testCLI.NewDatasetsCmd()
+	buf := new(bytes.Buffer)
+	cmd.SetOut(buf)
+
+	cmd.SetArgs([]string{datasetFile, privateKeyFile})
+	err = cmd.Execute()
+	require.NoError(t, err)
+
+	require.Contains(t, buf.String(), "Error decoding private key")
 	t.Cleanup(func() {
 		os.Remove(datasetFile)
 		os.Remove(privateKeyFile)
@@ -0,0 +1,46 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package cli
+
+import (
+	"github.com/absmach/magistrala/pkg/errors"
+	"github.com/fatih/color"
+	"github.com/spf13/cobra"
+	"github.com/ultravioletrs/cocos/agent/auth"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+var (
+	errAgentUnavailable                   = errors.New("agent is unavailable on the current address")
+	errDigitalSignatureVerificationFailed = errors.New("digital signature verification failed, check the provided public key")
+)
+
+func decodeErros(err error) error {
+	statusErr, ok := status.FromError(err)
+	if ok {
+		switch statusErr.Code() {
+		case codes.PermissionDenied:
+			return errDigitalSignatureVerificationFailed
+		case codes.Unavailable:
+			return errAgentUnavailable
+		case codes.Unknown:
+			return err
+		}
+	}
+	switch {
+	case errors.Contains(err, auth.ErrSignatureVerificationFailed):
+		return auth.ErrSignatureVerificationFailed
+
+	default:
+		return err
+	}
+}
+
+func printError(cmd *cobra.Command, message string, err error) {
+	if !Verbose {
+		err = decodeErros(err)
+	}
+	msg := color.New(color.FgRed).Sprintf(message, err)
+	cmd.Println(msg)
+}
@@ -0,0 +1,104 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package cli
+
+import (
+	"bytes"
+	"errors"
+	"testing"
+
+	mgerrors "github.com/absmach/magistrala/pkg/errors"
+	"github.com/fatih/color"
+	"github.com/spf13/cobra"
+	"github.com/ultravioletrs/cocos/agent/auth"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+func TestDecodeErros(t *testing.T) {
+	tests := []struct {
+		name     string
+		input    error
+		expected error
+	}{
+		{
+			name:     "Permission Denied",
+			input:    status.Error(codes.PermissionDenied, "permission denied"),
+			expected: errDigitalSignatureVerificationFailed,
+		},
+		{
+			name:     "Unavailable",
+			input:    status.Error(codes.Unavailable, "service unavailable"),
+			expected: errAgentUnavailable,
+		},
+		{
+			name:     "Unknown",
+			input:    status.Error(codes.Unknown, "unknown error"),
+			expected: status.Error(codes.Unknown, "unknown error"),
+		},
+		{
+			name:     "Signature Verification Failed",
+			input:    mgerrors.Wrap(auth.ErrSignatureVerificationFailed, errors.New("wrapped error")),
+			expected: auth.ErrSignatureVerificationFailed,
+		},
+		{
+			name:     "Other Error",
+			input:    errors.New("other error"),
+			expected: errors.New("other error"),
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := decodeErros(tt.input)
+			if result.Error() != tt.expected.Error() {
+				t.Errorf("decodeErros(%v) = %v, want %v", tt.input, result, tt.expected)
+			}
+		})
+	}
+}
+
+func TestPrintError(t *testing.T) {
+	// Save the original color.NoColor value and restore it after the test
+	origNoColor := color.NoColor
+	color.NoColor = true
+	defer func() { color.NoColor = origNoColor }()
+
+	tests := []struct {
+		name     string
+		message  string
+		err      error
+		verbose  bool
+		expected string
+	}{
+		{
+			name:     "Non-verbose mode",
+			message:  "Error: %s",
+			err:      status.Error(codes.PermissionDenied, "permission denied"),
+			verbose:  false,
+			expected: "Error: digital signature verification failed, check the provided public key\n",
+		},
+		{
+			name:     "Verbose mode",
+			message:  "Error: %s",
+			err:      status.Error(codes.PermissionDenied, "permission denied"),
+			verbose:  true,
+			expected: "Error: rpc error: code = PermissionDenied desc = permission denied\n",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			Verbose = tt.verbose
+			cmd := &cobra.Command{}
+			buf := new(bytes.Buffer)
+			cmd.SetOut(buf)
+
+			printError(cmd, tt.message, tt.err)
+
+			if got := buf.String(); got != tt.expected {
+				t.Errorf("printError() output = %q, want %q", got, tt.expected)
+			}
+		})
+	}
+}
@@ -10,9 +10,7 @@ import (
 	"crypto/rsa"
 	"crypto/x509"
 	"encoding/pem"
-	"log"
 	"os"
-	"reflect"

 	"github.com/spf13/cobra"
 )
@@ -44,48 +42,64 @@ func (cli *CLI) NewKeysCmd() *cobra.Command {
 			case ECDSA:
 				privEcdsaKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
 				if err != nil {
-					log.Fatalf("Error generating keys: %v", err)
+					printError(cmd, "Error generating keys: %v ❌ ", err)
+					return
 				}

 				pubKeyBytes, err := x509.MarshalPKIXPublicKey(&privEcdsaKey.PublicKey)
 				if err != nil {
-					log.Fatalf("Error marshalling public key: %v", err)
+					printError(cmd, "Error marshalling public key: %v ❌ ", err)
+					return
 				}

-				generateAndWriteKeys(privEcdsaKey, pubKeyBytes, ecdsaKeyType)
+				if err := generateAndWriteKeys(privEcdsaKey, pubKeyBytes, ecdsaKeyType); err != nil {
+					printError(cmd, "Error generating and writing keys: %v ❌ ", err)
+					return
+				}

 			case ED25519:
 				pubEd25519Key, privEd25519Key, err := ed25519.GenerateKey(rand.Reader)
 				if err != nil {
-					log.Fatalf("Error generating keys: %v", err)
+					printError(cmd, "Error generating keys: %v ❌ ", err)
+					return
 				}
 				pubKey, err := x509.MarshalPKIXPublicKey(pubEd25519Key)
 				if err != nil {
-					log.Fatalf("Error marshalling public key: %v", err)
+					printError(cmd, "Error marshalling public key: %v ❌ ", err)
+					return
+				}
+				if err := generateAndWriteKeys(privEd25519Key, pubKey, ed25519KeyType); err != nil {
+					printError(cmd, "Error generating and writing keys: %v ❌ ", err)
+					return
 				}
-				generateAndWriteKeys(privEd25519Key, pubKey, ed25519KeyType)

-			// Default to RSA
 			default:
 				privKey, err := rsa.GenerateKey(rand.Reader, keyBitSize)
 				if err != nil {
-					log.Fatalf("Error generating keys: %v", err)
+					printError(cmd, "Error generating keys: %v ❌ ", err)
+					return
 				}

 				pubKeyBytes, err := x509.MarshalPKIXPublicKey(&privKey.PublicKey)
 				if err != nil {
-					log.Fatalf("Error marshalling public key: %v", err)
+					printError(cmd, "Error marshalling public key: %v ❌ ", err)
+					return
+				}
+				if err := generateAndWriteKeys(privKey, pubKeyBytes, rsaKeyType); err != nil {
+					printError(cmd, "Error generating and writing keys: %v ❌ ", err)
+					return
 				}
-				generateAndWriteKeys(privKey, pubKeyBytes, rsaKeyType)
 			}
+
+			cmd.Printf("Successfully generated public/private key pair of type: %s", KeyType)
 		},
 	}
 }

-func generateAndWriteKeys(privKey interface{}, pubKeyBytes []byte, keyType string) {
+func generateAndWriteKeys(privKey interface{}, pubKeyBytes []byte, keyType string) error {
 	privFile, err := os.Create(privateKeyFile)
 	if err != nil {
-		log.Fatalf("Error creating private key file: %v", err)
+		return err
 	}
 	defer privFile.Close()

@@ -99,22 +113,19 @@ func generateAndWriteKeys(privKey interface{}, pubKeyBytes []byte, keyType strin
 		b, err = x509.MarshalPKCS8PrivateKey(privKey)
 	}
 	if err != nil {
-		log.Printf("Error marshalling private key: %v", err)
-		return
+		return err
 	}

 	if err := pem.Encode(privFile, &pem.Block{
 		Type:  keyType,
 		Bytes: b,
 	}); err != nil {
-		log.Printf("Error encoding private key: %v", err)
-		return
+		return err
 	}

 	pubFile, err := os.Create(publicKeyFile)
 	if err != nil {
-		log.Printf("Error creating public key file: %v", err)
-		return
+		return err
 	}
 	defer pubFile.Close()

@@ -122,9 +133,8 @@ func generateAndWriteKeys(privKey interface{}, pubKeyBytes []byte, keyType strin
 		Type:  publicKeyType,
 		Bytes: pubKeyBytes,
 	}); err != nil {
-		log.Printf("Error encoding public key: %v", err)
-		return
+		return err
 	}

-	log.Printf("Successfully generated public/private key pair of type: %s", reflect.TypeOf(privKey).String())
+	return nil
 }
@@ -4,7 +4,6 @@ package cli

 import (
 	"encoding/pem"
-	"log"
 	"os"

 	"github.com/fatih/color"
@@ -20,12 +19,11 @@ func (cli *CLI) NewResultsCmd() *cobra.Command {
 		Example: "result <private_key_file_path>",
 		Args:    cobra.ExactArgs(1),
 		Run: func(cmd *cobra.Command, args []string) {
-			log.Println("⏳ Retrieving computation result file")
+			cmd.Println("⏳ Retrieving computation result file")

 			privKeyFile, err := os.ReadFile(args[0])
 			if err != nil {
-				msg := color.New(color.FgRed).Sprintf("Error reading private key file: %v ❌ ", err)
-				log.Println(msg)
+				printError(cmd, "Error reading private key file: %v ❌ ", err)
 				return
 			}

@@ -33,21 +31,23 @@ func (cli *CLI) NewResultsCmd() *cobra.Command {

 			var result []byte

-			privKey := decodeKey(pemBlock)
+			privKey, err := decodeKey(pemBlock)
+			if err != nil {
+				printError(cmd, "Error decoding private key: %v ❌ ", err)
+				return
+			}
 			result, err = cli.agentSDK.Result(cmd.Context(), privKey)
 			if err != nil {
-				msg := color.New(color.FgRed).Sprintf("Error retrieving computation result: %v ❌ ", err)
-				log.Println(msg)
+				printError(cmd, "Error retrieving computation result: %v ❌ ", err)
 				return
 			}

 			if err := os.WriteFile(resultFilePath, result, 0o644); err != nil {
-				msg := color.New(color.FgRed).Sprintf("Error saving computation result to %s: %v  ❌ ", resultFilePath, err)
-				log.Println(msg)
+				printError(cmd, "Error saving computation result file: %v  ❌ ", err)
 				return
 			}

-			log.Println(color.New(color.FgGreen).Sprint("Computation result retrieved and saved successfully! ✔ "))
+			cmd.Println(color.New(color.FgGreen).Sprint("Computation result retrieved and saved successfully! ✔ "))
 		},
 	}
 }
@@ -3,6 +3,7 @@
 package cli

 import (
+	"bytes"
 	"errors"
 	"os"
 	"testing"
@@ -23,13 +24,13 @@ func TestResultsCmd_Success(t *testing.T) {
 	require.NoError(t, err)

 	cmd := testCLI.NewResultsCmd()
-	output := captureLogOutput(func() {
-		cmd.SetArgs([]string{privateKeyFile})
-		err = cmd.Execute()
-		require.NoError(t, err)
-	})
+	buf := new(bytes.Buffer)
+	cmd.SetOut(buf)
+	cmd.SetArgs([]string{privateKeyFile})
+	err = cmd.Execute()
+	require.NoError(t, err)

-	require.Contains(t, output, "Computation result retrieved and saved successfully")
+	require.Contains(t, buf.String(), "Computation result retrieved and saved successfully")

 	resultFile, err := os.ReadFile("results.zip")
 	require.NoError(t, err)
@@ -47,13 +48,13 @@ func TestResultsCmd_MissingPrivateKeyFile(t *testing.T) {
 	testCLI := New(mockSDK)

 	cmd := testCLI.NewResultsCmd()
-	output := captureLogOutput(func() {
-		cmd.SetArgs([]string{"non_existent_private_key.pem"})
-		err := cmd.Execute()
-		require.NoError(t, err)
-	})
+	buf := new(bytes.Buffer)
+	cmd.SetOut(buf)
+	cmd.SetArgs([]string{"non_existent_private_key.pem"})
+	err := cmd.Execute()
+	require.NoError(t, err)

-	require.Contains(t, output, "Error reading private key file")
+	require.Contains(t, buf.String(), "Error reading private key file")
 }

 func TestResultsCmd_ResultFailure(t *testing.T) {
@@ -65,13 +66,13 @@ func TestResultsCmd_ResultFailure(t *testing.T) {
 	require.NoError(t, err)

 	cmd := testCLI.NewResultsCmd()
-	output := captureLogOutput(func() {
-		cmd.SetArgs([]string{privateKeyFile})
-		err = cmd.Execute()
-		require.NoError(t, err)
-	})
+	buf := new(bytes.Buffer)
+	cmd.SetOut(buf)
+	cmd.SetArgs([]string{privateKeyFile})
+	err = cmd.Execute()
+	require.NoError(t, err)

-	require.Contains(t, output, "error retrieving computation result")
+	require.Contains(t, buf.String(), "error retrieving computation result")
 	mockSDK.AssertCalled(t, "Result", mock.Anything, mock.Anything)
 	t.Cleanup(func() {
 		os.Remove(privateKeyFile)
@@ -91,13 +92,13 @@ func TestResultsCmd_SaveFailure(t *testing.T) {
 	require.NoError(t, err)

 	cmd := testCLI.NewResultsCmd()
-	output := captureLogOutput(func() {
-		cmd.SetArgs([]string{privateKeyFile})
-		err := cmd.Execute()
-		require.NoError(t, err)
-	})
+	buf := new(bytes.Buffer)
+	cmd.SetOut(buf)
+	cmd.SetArgs([]string{privateKeyFile})
+	err = cmd.Execute()
+	require.NoError(t, err)

-	require.Contains(t, output, "Error saving computation result to results.zip")
+	require.Contains(t, buf.String(), "Error saving computation result file")
 	mockSDK.AssertCalled(t, "Result", mock.Anything, mock.Anything)

 	t.Cleanup(func() {
@@ -105,3 +106,29 @@ func TestResultsCmd_SaveFailure(t *testing.T) {
 		os.Remove(privateKeyFile)
 	})
 }
+
+func TestResultsCmd_InvalidPrivateKey(t *testing.T) {
+	mockSDK := new(mocks.SDK)
+	mockSDK.On("Result", mock.Anything, mock.Anything).Return([]byte(compResult), nil)
+	testCLI := New(mockSDK)
+
+	invalidPrivateKey, err := os.CreateTemp("", "invalid_private_key.pem")
+	require.NoError(t, err)
+	err = invalidPrivateKey.Close()
+	require.NoError(t, err)
+
+	t.Cleanup(func() {
+		err := os.Remove(invalidPrivateKey.Name())
+		require.NoError(t, err)
+	})
+
+	cmd := testCLI.NewResultsCmd()
+	buf := new(bytes.Buffer)
+	cmd.SetOut(buf)
+	cmd.SetArgs([]string{invalidPrivateKey.Name()})
+	err = cmd.Execute()
+	require.NoError(t, err)
+
+	require.Contains(t, buf.String(), "Error decoding private key")
+	mockSDK.AssertNotCalled(t, "Result", mock.Anything, mock.Anything)
+}
@@ -4,6 +4,8 @@ package cli

 import "github.com/ultravioletrs/cocos/pkg/sdk"

+var Verbose bool
+
 type CLI struct {
 	agentSDK sdk.SDK
 }