support multipart encyrption upload

Signed-off-by: Dweb Fan <[email protected]>

cmd/lomob/crypt.go

@@ -68,10 +68,6 @@ func encryptCmd(ctx *cli.Context) error {
 			return err
 		}
 	}
-	encryptKey, iv, err := genEncryptKeyAndSalt([]byte(masterKey))
-	if err != nil {
-		return err
-	}
 
 	src, err := os.Open(ifilename)
 	if err != nil {
@@ -85,13 +81,8 @@ func encryptCmd(ctx *cli.Context) error {
 	}
 	defer dst.Close()
 
-	encryptor, err := crypto.NewEncryptor(src, encryptKey, iv)
-	if err != nil {
-		return err
-	}
-
 	fmt.Printf("Start encrypt '%s', and save output to '%s'\n", ifilename, ofilename)
-	_, err = io.Copy(dst, encryptor)
+	_, err = encryptLocalFile(src, dst, masterKey, true)
 	if err != nil {
 		return err
 	}
@@ -101,6 +92,24 @@ func encryptCmd(ctx *cli.Context) error {
 	return nil
 }
 
+func encryptLocalFile(src io.ReadSeeker, dst io.Writer, masterKey string, hasHeader bool) ([]byte, error) {
+	encryptKey, iv, err := genEncryptKeyAndSalt([]byte(masterKey))
+	if err != nil {
+		return nil, err
+	}
+
+	encryptor, err := crypto.NewEncryptor(src, encryptKey, iv, hasHeader)
+	if err != nil {
+		return nil, err
+	}
+
+	_, err = io.Copy(dst, encryptor)
+	if err != nil {
+		return nil, err
+	}
+	return encryptor.GetHash(), nil
+}
+
 func decryptLocalFile(ctx *cli.Context) error {
 	if len(ctx.Args()) != 1 {
 		return errors.New("usage: [encrypted file name]")

cmd/lomob/main.go

@@ -132,6 +132,24 @@ func main() {
 							Name:  "save-parts,s",
 							Usage: "Save multiparts locally for debug",
 						},
+						cli.BoolFlag{
+							Name:  "no-encrypt",
+							Usage: "not do any encryption, and upload raw files",
+						},
+						cli.BoolFlag{
+							Name:  "force",
+							Usage: "force to upload from scratch and not reuse previous upload info",
+						},
+						cli.StringFlag{
+							Name:   "encrypt-key, k",
+							Usage:  "Master key to encrypt current upload file",
+							EnvVar: "LOMOB_MASTER_KEY",
+						},
+						cli.StringFlag{
+							Name:  "storage-class",
+							Usage: "The  type  of storage to use for the object. Valid choices are: DEEP_ARCHIVE | GLACIER | GLACIER_IR | INTELLIGENT_TIERING | ONE-ZONE_IA | REDUCED_REDUNDANCY | STANDARD | STANDARD_IA.",
+							Value: "STANDARD",
+						},
 					},
 				},
 			},
@@ -179,6 +197,14 @@ func main() {
 							Name:  "save-parts,s",
 							Usage: "Save multiparts locally for debug",
 						},
+						cli.BoolFlag{
+							Name:  "no-encrypt",
+							Usage: "not do any encryption, and upload raw files",
+						},
+						cli.BoolFlag{
+							Name:  "force",
+							Usage: "force to upload from scratch and not reuse previous upload info",
+						},
 						cli.StringFlag{
 							Name:   "encrypt-key, k",
 							Usage:  "Master key to encrypt current upload file",
@@ -187,13 +213,13 @@ func main() {
 						cli.StringFlag{
 							Name:  "storage-class",
 							Usage: "The  type  of storage to use for the object. Valid choices are: DEEP_ARCHIVE | GLACIER | GLACIER_IR | INTELLIGENT_TIERING | ONE-ZONE_IA | REDUCED_REDUNDANCY | STANDARD | STANDARD_IA.",
-							Value: "GLACIER_IR",
+							Value: "STANDARD",
 						},
 					},
 				},
 				{
 					Name:   "files",
-					Action: uploadFiles,
+					Action: uploadFilesToGdrive,
 					Usage:  "Upload individual files not in ISO to google drive",
 					Flags: []cli.Flag{
 						cli.StringFlag{
@@ -461,7 +487,7 @@ func main() {
 				},
 				{
 					Name:      "upload-s3",
-					Action:    uploadFileToS3,
+					Action:    uploadFilesToS3,
 					Usage:     "Upload individual file into S3 with on-the-fly encryption",
 					ArgsUsage: "[local file name]",
 					Flags: []cli.Flag{
@@ -485,15 +511,19 @@ func main() {
 							Usage: "awsBucketName",
 							Value: defaultBucket,
 						},
+						cli.BoolFlag{
+							Name:  "no-encrypt",
+							Usage: "not do any encryption, and upload raw files",
+						},
 						cli.StringFlag{
 							Name:   "encrypt-key, k",
-							Usage:  "Master key to encrypt current upload file",
+							Usage:  "Master key to encrypt current upload file. If it is empty, means no encryption is needed",
 							EnvVar: "LOMOB_MASTER_KEY",
 						},
 						cli.StringFlag{
 							Name:  "storage-class",
 							Usage: "The  type  of storage to use for the object. Valid choices are: DEEP_ARCHIVE | GLACIER | GLACIER_IR | INTELLIGENT_TIERING | ONE-ZONE_IA | REDUCED_REDUNDANCY | STANDARD | STANDARD_IA.",
-							Value: "GLACIER_IR",
+							Value: "STANDARD",
 						},
 					},
 				},

cmd/lomob/upload-files.go

@@ -18,7 +18,7 @@ import (
 	"github.com/urfave/cli"
 )
 
-func uploadFiles(ctx *cli.Context) error {
+func uploadFilesToGdrive(ctx *cli.Context) error {
 	err := initDB(ctx.GlobalString("db"))
 	if err != nil {
 		return err
@@ -153,7 +153,7 @@ func uploadFiles(ctx *cli.Context) error {
 			return err
 		}
 
-		encryptor, err := crypto.NewEncryptor(file, encryptKey, iv)
+		encryptor, err := crypto.NewEncryptor(file, encryptKey, iv, true)
 		if err != nil {
 			return err
 		}
@@ -195,7 +195,7 @@ func flattenScanRootDir(dir string) string {
 	return strings.ReplaceAll(dir, string(os.PathSeparator), "_")
 }
 
-func uploadFileToS3(ctx *cli.Context) error {
+func uploadFilesToS3(ctx *cli.Context) error {
 	accessKeyID := ctx.String("awsAccessKeyID")
 	accessKey := ctx.String("awsSecretAccessKey")
 	region := ctx.String("awsBucketRegion")
@@ -206,20 +206,45 @@ func uploadFileToS3(ctx *cli.Context) error {
 		return err
 	}
 
+	masterKey := ctx.String("encrypt-key")
+	if ctx.Bool("no-encrypt") {
+		masterKey = ""
+	} else if masterKey == "" {
+		masterKey, err = getMasterKey()
+		if err != nil {
+			return err
+		}
+	}
+
 	cli, err := clients.NewAWSClient(accessKeyID, accessKey, region)
 	if err != nil {
 		return err
 	}
 
-	masterKey := ctx.String("encrypt-key")
-	if masterKey == "" {
-		masterKey, err = getMasterKey()
+	for _, name := range ctx.Args() {
+		fmt.Printf("Uploading file %s\n", name)
+
+		if masterKey == "" {
+			err = uploadRawFileToS3(cli, bucket, storageClass, name, binContentType)
+			if err != nil {
+				return err
+			}
+			continue
+		}
+		tmpFilename, err := uploadEncryptFileToS3(cli, bucket, storageClass, name, masterKey)
+		if err != nil {
+			return err
+		}
+		err = os.Remove(tmpFilename)
 		if err != nil {
 			return err
 		}
 	}
+	return nil
+}
 
-	remoteFilename := filepath.Base(ctx.Args()[0])
+func uploadFileToS3(cli *clients.AWSClient, bucket, storageClass, remoteFilename, expectHash, contentType string, expectSize int,
+	reader io.ReadSeeker) error {
 	remoteInfo, err := cli.HeadObject(bucket, remoteFilename)
 	if err != nil {
 		return err
@@ -229,53 +254,88 @@ func uploadFileToS3(ctx *cli.Context) error {
 			remoteFilename, bucket)
 		return nil
 	}
+	if remoteInfo != nil {
+		recreate := false
+		if remoteInfo.Size != expectSize {
+			logrus.Warnf("%s exists in cloud and its size is %d, but provided file size is %d",
+				remoteFilename, remoteInfo.Size, expectSize)
+			recreate = true
+		}
+		if remoteInfo.HashBase64 != expectHash {
+			logrus.Warnf("%s exists in cloud and its checksum is %s, but provided checksum is %s",
+				remoteFilename, remoteInfo.HashBase64, expectHash)
+			recreate = true
+		}
+		// no need upload, return nil upload request
+		if !recreate {
+			fmt.Printf("%s is already in bucket %s, no need upload again !\n",
+				remoteFilename, bucket)
+			return nil
+		}
+	}
 
-	src, err := os.Open(ctx.Args()[0])
+	err = cli.PutObject(bucket, remoteFilename, expectHash, contentType, storageClass, reader)
 	if err != nil {
-		return err
+		fmt.Printf("Uploading metadata file %s fail: %s\n", remoteFilename, err)
+	} else {
+		fmt.Printf("Upload metadata file %s success!\n", remoteFilename)
 	}
-	defer src.Close()
+	return err
+}
 
-	encryptKey, iv, err := genEncryptKeyAndSalt([]byte(masterKey))
+func uploadRawFileToS3(cli *clients.AWSClient, bucket, storageClass, filename, contentType string) error {
+	h, err := hash.CalculateHashFile(filename)
 	if err != nil {
 		return err
 	}
 
-	encryptor, err := crypto.NewEncryptor(src, encryptKey, iv)
+	hashBase64 := hash.CalculateHashBase64(h)
+
+	f, err := os.Open(filename)
 	if err != nil {
 		return err
 	}
+	defer f.Close()
 
-	// as PutObject requires encryption before input, thus, it has to write into one temp file
-	tmpFile, err := os.CreateTemp("", "")
+	stat, err := f.Stat()
 	if err != nil {
 		return err
 	}
-	tmpFileName := tmpFile.Name()
-	defer os.Remove(tmpFileName)
 
-	_, err = io.Copy(tmpFile, encryptor)
+	return uploadFileToS3(cli, bucket, storageClass, filepath.Base(filename), hashBase64, contentType, int(stat.Size()), f)
+}
+
+// as PutObject requires encryption before input, thus, it has to write into one temp file or memory to get all data
+// return tmp filename, and let caller delete
+func uploadEncryptFileToS3(cli *clients.AWSClient, bucket, storageClass, filename, masterKey string) (string, error) {
+	src, err := os.Open(filename)
 	if err != nil {
-		return err
+		return "", err
 	}
+	defer src.Close()
 
-	hash, err := lomohash.CalculateHashFile(tmpFileName)
+	tmpFile, err := os.CreateTemp("", "")
 	if err != nil {
-		return err
+		return "", err
 	}
+	tmpFileName := tmpFile.Name()
+	defer tmpFile.Close()
 
-	_, err = tmpFile.Seek(0, io.SeekStart)
+	hash, err := encryptLocalFile(src, tmpFile, masterKey, true)
 	if err != nil {
-		return err
+		return "", err
 	}
 
-	fmt.Printf("Uploading file %s\n", remoteFilename)
-	err = cli.PutObject(bucket, remoteFilename, lomohash.CalculateHashBase64(hash), metaContentType, storageClass, tmpFile)
+	size, err := tmpFile.Seek(0, io.SeekEnd)
 	if err != nil {
-		fmt.Printf("Uploading file %s fail: %s\n", remoteFilename, err)
-	} else {
-		fmt.Printf("Upload file %s success!\n", remoteFilename)
+		return "", err
 	}
 
-	return err
+	_, err = tmpFile.Seek(0, io.SeekStart)
+	if err != nil {
+		return "", err
+	}
+
+	return tmpFileName, uploadFileToS3(cli, bucket, storageClass, filepath.Base(filename), lomohash.CalculateHashBase64(hash),
+		binContentType, int(size), tmpFile)
 }