added some update

Abdiooa · Feb 28, 2024 · 911ddbf · 911ddbf
1 parent 3a5c1c5
commit 911ddbf
Show file tree

Hide file tree

Showing 7 changed files with 160 additions and 135 deletions.
diff --git a/cmd/zeroward/download.go b/cmd/zeroward/download.go
@@ -6,76 +6,52 @@ import (
 	"github.com/Abdiooa/zeroward/pkg/zeroward/common"
 	"github.com/Abdiooa/zeroward/pkg/zeroward/downloading"
 	"github.com/spf13/cobra"
-	"github.com/spf13/viper"
 )
 
 // downloadCmd represents the download command
 var downloadCmd = &cobra.Command{
 	Use:   "download",
 	Short: "Download Command to download Files from the cloud.",
-	Long:  `This command is used to download User Objects(Files) from the cloud storage.`,
+	Long:  `This command is used to download User Objects(Files) from the cloud storage, also gives you an readable files if the files were encrypted.`,
 	Run: func(cmd *cobra.Command, args []string) {
-
 		accessKeyID, _ := cmd.Flags().GetString("accessKeyID")
-
 		secretAccessKey, _ := cmd.Flags().GetString("secretAccessKey")
-
-		// passphrase, _ := cmd.Flags().GetString("passphrase")
-
-		bcktName, _ := cmd.Flags().GetString("bcktname")
-
+		bcktName, _ := cmd.Flags().GetString("bucketName")
 		filePath, _ := cmd.Flags().GetString("filePath")
-
 		objectkey, _ := cmd.Flags().GetString("objectkey")
+		removeAfterDownload, _ := cmd.Flags().GetBool("remove")
+		decryptWhileDownloading, _ := cmd.Flags().GetBool("decrypt")
 
-		removeAfterDownload, _ := cmd.Flags().GetString("removeAfterDownload")
-
-		KeyAccessDefined := common.IsNotKeyAccessDefined()
-		region := viper.GetString("Region")
-		if KeyAccessDefined {
-
-			if accessKeyID == "" || secretAccessKey == "" || bcktName == "" {
-				fmt.Println("Error: Access Key ID and Secret Access Key are required as for your login and password of the Cloud Storage, also the bucket name is required!")
-				return
-			}
-
-			if filePath == "" || objectkey == "" {
-				fmt.Println("Error: FilePath and ObjectKey are required!")
+		if decryptWhileDownloading {
+			if err := handleDownload(accessKeyID, secretAccessKey, bcktName, filePath, objectkey, removeAfterDownload); err != nil {
+				fmt.Println("Error:", err)
 				return
 			}
-
-			common.StoreCredentials(accessKeyID, secretAccessKey, region)
-
-			err := downloading.DownloadObject(region, accessKeyID, secretAccessKey, bcktName, filePath, objectkey, removeAfterDownload)
-
+		} else {
+			accessKeyID, secretAccessKey, region, err := common.HandleCredentials(accessKeyID, secretAccessKey)
 			if err != nil {
 				fmt.Println("Error:", err)
 				return
 			}
-		}
-
-		if bcktName == "" {
-			fmt.Println("Error:  the bucket name is required!")
-			return
-		}
-
-		if !KeyAccessDefined {
-
-			if filePath == "" || objectkey == "" {
-				fmt.Println("Error: FilePath and ObjectKey are required!")
+			if err := downloading.DownloadNormalObject(region, accessKeyID, secretAccessKey, bcktName, filePath, objectkey, removeAfterDownload); err != nil {
+				fmt.Println("Error:", err)
 				return
-			} else {
-				err := downloading.DownloadObject(region, viper.GetString("AWSAccessKeyID"), viper.GetString("AWSSecretAccessKey"), bcktName, filePath, objectkey, removeAfterDownload)
-				if err != nil {
-					fmt.Println("Error:", err)
-					return
-				}
 			}
 		}
 	},
 }
 
+func handleDownload(accessKeyID, secretAccessKey, bcktName, filePath, objectkey string, removeAfterDownload bool) error {
+	accessKeyID, secretAccessKey, region, err := common.HandleCredentials(accessKeyID, secretAccessKey)
+	if err != nil {
+		return err
+	}
+
+	return downloading.DownloadObject(region, accessKeyID, secretAccessKey, bcktName, filePath, objectkey, removeAfterDownload)
+}
+
 func init() {
 	rootCmd.AddCommand(downloadCmd)
-	downloadCmd.Flags().StringP("removeAfterDownload", "r", "", "write (yes/y) for removing the file from the cloud storage after successful download, else don't specify anything")
+	downloadCmd.Flags().BoolP("remove", "r", false, "Remove the file from the cloud storage after successful download")
+	downloadCmd.Flags().BoolP("decrypt", "d", false, "Decrypt file while downloading")
 }
diff --git a/cmd/zeroward/objects.go b/cmd/zeroward/objects.go
@@ -20,7 +20,7 @@ var ObjectsCmd = &cobra.Command{
 
 		secretAccessKey, _ := cmd.Flags().GetString("secretAccessKey")
 
-		bcktName, _ := cmd.Flags().GetString("bcktname")
+		bcktName, _ := cmd.Flags().GetString("bucketName")
 
 		KeyAccessDefined := common.IsNotKeyAccessDefined()
 		region := viper.GetString("Region")

diff --git a/cmd/zeroward/removeFile.go b/cmd/zeroward/removeFile.go
@@ -20,7 +20,7 @@ var removeCmd = &cobra.Command{
 
 		secretAccessKey, _ := cmd.Flags().GetString("secretAccessKey")
 
-		bcktName, _ := cmd.Flags().GetString("bcktname")
+		bcktName, _ := cmd.Flags().GetString("bucketName")
 
 		objectkey, _ := cmd.Flags().GetString("objectkey")
 

diff --git a/cmd/zeroward/root.go b/cmd/zeroward/root.go
@@ -66,7 +66,7 @@ func init() {
 	rootCmd.PersistentFlags().StringP("filePath", "f", "", "Path of the file that you want to encrypt")
 	rootCmd.PersistentFlags().StringP("accessKeyID", "i", "", "Access Key Id as your Login Key")
 	rootCmd.PersistentFlags().StringP("secretAccessKey", "s", "", "Secret Access Key as your Password")
-	rootCmd.PersistentFlags().StringP("bcktname", "b", "", "Bucket name out of all your existing buckets!")
+	rootCmd.PersistentFlags().StringP("bucketName", "b", "", "Bucket name out of all your existing buckets!")
 	rootCmd.PersistentFlags().StringP("objectkey", "o", "", "objectkey refers to the unique identifier or name of the object(file) with a bucket, it is the path or where in the bucket the file should be stored.")
 	rootCmd.PersistentFlags().StringP("passphrase", "p", "", "Passphrase for encryption required for the first encryption")
 

diff --git a/cmd/zeroward/upload.go b/cmd/zeroward/upload.go
@@ -14,120 +14,85 @@ import (
 
 const DEKKeyMetadataKey = "dek-key" // Metadata key for the DEK key
 
-// uploadCmd represents the upload command
 var uploadCmd = &cobra.Command{
 	Use:   "upload",
 	Short: "Upload Command to upload Files on a cloud",
 	Long:  `This Command is used to upload a secured(encrypted) user file on the cloud storage.`,
 	Run: func(cmd *cobra.Command, args []string) {
-
 		accessKeyID, _ := cmd.Flags().GetString("accessKeyID")
-
 		secretAccessKey, _ := cmd.Flags().GetString("secretAccessKey")
-
-		bcktName, _ := cmd.Flags().GetString("bcktname")
-
+		bcktName, _ := cmd.Flags().GetString("bucketName")
 		passphrase, _ := cmd.Flags().GetString("passphrase")
-
 		filePath, _ := cmd.Flags().GetString("filePath")
-
 		objectkey, _ := cmd.Flags().GetString("objectkey")
+		encrypt, _ := cmd.Flags().GetBool("encrypt") // Change to GetBool
 
-		firstEncryption := common.IsFirstEncryption()
-
-		if firstEncryption {
-			if passphrase == "" {
-				fmt.Println("Error: Passphrase is required for the first encryption. Please provide a passphrase using the --passphrase flag.")
+		if encrypt {
+			if err := handleEncryptionAndUpload(accessKeyID, secretAccessKey, bcktName, filePath, passphrase, objectkey); err != nil {
+				fmt.Println("Error:", err)
 				return
 			}
-
-			kekKey, err := genekeys.GenerateKek(passphrase)
-			if err != nil {
+		} else {
+			if err := handleUpload(accessKeyID, secretAccessKey, bcktName, filePath, objectkey, nil); err != nil {
 				fmt.Println("Error:", err)
 				return
 			}
-			common.UpdateKEKKey(kekKey)
 		}
+	},
+}
 
-		kekk := viper.GetString("KEKkey")
-
-		kekBytes, err := hex.DecodeString(kekk)
-
-		if err != nil {
-			fmt.Println("Error:", err)
-			return
+func handleEncryptionAndUpload(accessKeyID, secretAccessKey, bcktName, filePath, passphrase, objectkey string) error {
+	if common.IsFirstEncryption() {
+		if passphrase == "" {
+			return fmt.Errorf("error: Passphrase is required for the first encryption. Please provide a passphrase using the --passphrase flag")
 		}
-
-		dek, err := genekeys.GenerateDek()
-		cobra.CheckErr(err)
-
-		if filePath != "" {
-			if err := encryption.EncryptFile(filePath, dek); err != nil {
-				fmt.Println("Error encrypting File:", err)
-				return
-			}
-		}
-		encryptedDek, err := encryption.EncryptData(dek, kekBytes)
-
+		kekKey, err := genekeys.GenerateKek(passphrase)
 		if err != nil {
-			fmt.Println("Error:", err)
-			return
+			return fmt.Errorf("error generating KEK key: %v", err)
 		}
-
-		dekString := hex.EncodeToString(encryptedDek)
-
-		// Create metadata map and add DEK key
-		metadata := map[string]string{
-			DEKKeyMetadataKey: dekString,
+		common.UpdateKEKKey(kekKey)
+	}
+
+	kekk, err := hex.DecodeString(viper.GetString("KEKkey"))
+	if err != nil {
+		return fmt.Errorf("error decoding KEK key: %v", err)
+	}
+
+	dek, err := genekeys.GenerateDek()
+	if err != nil {
+		return fmt.Errorf("error generating DEK: %v", err)
+	}
+
+	if filePath != "" {
+		if err := encryption.EncryptFile(filePath, dek); err != nil {
+			return fmt.Errorf("error encrypting file: %v", err)
 		}
+	}
 
-		KeyAccessDefined := common.IsNotKeyAccessDefined()
-		region := viper.GetString("Region")
+	encryptedDek, err := encryption.EncryptData(dek, kekk)
+	if err != nil {
+		return fmt.Errorf("error encrypting DEK: %v", err)
+	}
 
-		if KeyAccessDefined {
-			if accessKeyID == "" || secretAccessKey == "" || bcktName == "" {
-				fmt.Println("Error: Access Key ID and Secret Access Key are required as for your login and password of the Cloud Storage, also the bucket name is required!")
-				return
-			}
-			if filePath == "" {
-				fmt.Println("Error: FilePath are required!")
-				return
-			}
-			common.StoreCredentials(accessKeyID, secretAccessKey, region)
+	dekString := hex.EncodeToString(encryptedDek)
 
-			outputFilePath := fmt.Sprintf("%s.enc", filePath)
-			err := uploading.UploadFile(region, accessKeyID, secretAccessKey, bcktName, outputFilePath, objectkey, metadata)
+	metadata := map[string]string{
+		DEKKeyMetadataKey: dekString,
+	}
 
-			if err != nil {
-				fmt.Println("Error:", err)
-				return
-			}
-		}
+	return handleUpload(accessKeyID, secretAccessKey, bcktName, fmt.Sprintf("%s.enc", filePath), objectkey, metadata)
+}
 
-		if bcktName == "" {
-			fmt.Println("Error:  the bucket name is required!")
-			return
-		}
+func handleUpload(accessKeyID, secretAccessKey, bcktName, filePath, objectkey string, metadata map[string]string) error {
+	accessKeyID, secretAccessKey, region, err := common.HandleCredentials(accessKeyID, secretAccessKey)
+	if err != nil {
+		return err
+	}
 
-		if !KeyAccessDefined {
-			if bcktName != "" {
-
-				if filePath == "" {
-					fmt.Println("Error: FilePath is required!")
-					return
-				} else {
-					outputFilePath := fmt.Sprintf("%s.enc", filePath)
-					err := uploading.UploadFile(region, viper.GetString("AWSAccessKeyID"), viper.GetString("AWSSecretAccessKey"), bcktName, outputFilePath, objectkey, metadata)
-					if err != nil {
-						fmt.Println("Error:", err)
-						return
-					}
-				}
-			}
-		}
-	},
+	return uploading.UploadFile(region, accessKeyID, secretAccessKey, bcktName, filePath, objectkey, metadata)
 }
 
 func init() {
 	rootCmd.AddCommand(uploadCmd)
+	uploadCmd.Flags().BoolP("encrypt", "e", false, "Encrypt file before uploading")
 }
diff --git a/pkg/zeroward/common/credentialscrea.go b/pkg/zeroward/common/credentialscrea.go
@@ -94,3 +94,19 @@ func UpdateKEKKey(kek []byte) {
 		os.Exit(1)
 	}
 }
+
+func HandleCredentials(accessKeyID, secretAccessKey string) (string, string, string, error) {
+	KeyAccessDefined := IsNotKeyAccessDefined()
+	region := viper.GetString("Region")
+
+	if KeyAccessDefined {
+		if accessKeyID == "" || secretAccessKey == "" {
+			return "", "", "", fmt.Errorf("error: Access Key ID and Secret Access Key are required as for your login and password of the Cloud Storage, also the bucket name is required")
+		}
+
+		StoreCredentials(accessKeyID, secretAccessKey, region)
+		return accessKeyID, secretAccessKey, region, nil
+	}
+
+	return viper.GetString("AWSAccessKeyID"), viper.GetString("AWSSecretAccessKey"), region, nil
+}