github.com/iikira/iikira-go-utils@v0.0.0-20230610031953-f2cb11cde33a/utils/crypto.go (about)

     1  package utils
     2  
     3  import (
     4  	"fmt"
     5  	"github.com/iikira/iikira-go-utils/utils/bdcrypto"
     6  	"io"
     7  	"os"
     8  	"strings"
     9  )
    10  
    11  // CryptoMethodSupport 检测是否支持加密解密方法
    12  func CryptoMethodSupport(method string) bool {
    13  	switch method {
    14  	case "aes-128-ctr", "aes-192-ctr", "aes-256-ctr", "aes-128-cfb", "aes-192-cfb", "aes-256-cfb", "aes-128-ofb", "aes-192-ofb", "aes-256-ofb":
    15  		return true
    16  	}
    17  
    18  	return false
    19  }
    20  
    21  // EncryptFile 加密本地文件
    22  func EncryptFile(method string, key []byte, filePath string, isGzip bool) (encryptedFilePath string, err error) {
    23  	if !CryptoMethodSupport(method) {
    24  		return "", fmt.Errorf("unknown encrypt method: %s", method)
    25  	}
    26  
    27  	if isGzip {
    28  		err = bdcrypto.GZIPCompressFile(filePath)
    29  		if err != nil {
    30  			return
    31  		}
    32  	}
    33  
    34  	plainFile, err := os.OpenFile(filePath, os.O_RDONLY, 0)
    35  	if err != nil {
    36  		return
    37  	}
    38  
    39  	defer plainFile.Close()
    40  
    41  	var cipherReader io.Reader
    42  	switch method {
    43  	case "aes-128-ctr":
    44  		cipherReader, err = bdcrypto.Aes128CTREncrypt(bdcrypto.Convert16bytes(key), plainFile)
    45  	case "aes-192-ctr":
    46  		cipherReader, err = bdcrypto.Aes192CTREncrypt(bdcrypto.Convert24bytes(key), plainFile)
    47  	case "aes-256-ctr":
    48  		cipherReader, err = bdcrypto.Aes256CTREncrypt(bdcrypto.Convert32bytes(key), plainFile)
    49  	case "aes-128-cfb":
    50  		cipherReader, err = bdcrypto.Aes128CFBEncrypt(bdcrypto.Convert16bytes(key), plainFile)
    51  	case "aes-192-cfb":
    52  		cipherReader, err = bdcrypto.Aes192CFBEncrypt(bdcrypto.Convert24bytes(key), plainFile)
    53  	case "aes-256-cfb":
    54  		cipherReader, err = bdcrypto.Aes256CFBEncrypt(bdcrypto.Convert32bytes(key), plainFile)
    55  	case "aes-128-ofb":
    56  		cipherReader, err = bdcrypto.Aes128OFBEncrypt(bdcrypto.Convert16bytes(key), plainFile)
    57  	case "aes-192-ofb":
    58  		cipherReader, err = bdcrypto.Aes192OFBEncrypt(bdcrypto.Convert24bytes(key), plainFile)
    59  	case "aes-256-ofb":
    60  		cipherReader, err = bdcrypto.Aes256OFBEncrypt(bdcrypto.Convert32bytes(key), plainFile)
    61  	default:
    62  		return "", fmt.Errorf("unknown encrypt method: %s", method)
    63  	}
    64  
    65  	if err != nil {
    66  		return
    67  	}
    68  
    69  	plainFileInfo, err := plainFile.Stat()
    70  	if err != nil {
    71  		return
    72  	}
    73  
    74  	encryptedFilePath = filePath + ".encrypt"
    75  	encryptedFile, err := os.OpenFile(encryptedFilePath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, plainFileInfo.Mode())
    76  	if err != nil {
    77  		return
    78  	}
    79  
    80  	defer encryptedFile.Close()
    81  
    82  	_, err = io.Copy(encryptedFile, cipherReader)
    83  	if err != nil {
    84  		return
    85  	}
    86  
    87  	os.Remove(filePath)
    88  
    89  	return encryptedFilePath, nil
    90  }
    91  
    92  // DecryptFile 加密本地文件
    93  func DecryptFile(method string, key []byte, filePath string, isGzip bool) (decryptedFilePath string, err error) {
    94  	if !CryptoMethodSupport(method) {
    95  		return "", fmt.Errorf("unknown decrypt method: %s", method)
    96  	}
    97  
    98  	cipherFile, err := os.OpenFile(filePath, os.O_RDONLY, 0644)
    99  	if err != nil {
   100  		return
   101  	}
   102  
   103  	defer cipherFile.Close()
   104  
   105  	var plainReader io.Reader
   106  	switch method {
   107  	case "aes-128-ctr":
   108  		plainReader, err = bdcrypto.Aes128CTRDecrypt(bdcrypto.Convert16bytes(key), cipherFile)
   109  	case "aes-192-ctr":
   110  		plainReader, err = bdcrypto.Aes192CTRDecrypt(bdcrypto.Convert24bytes(key), cipherFile)
   111  	case "aes-256-ctr":
   112  		plainReader, err = bdcrypto.Aes256CTRDecrypt(bdcrypto.Convert32bytes(key), cipherFile)
   113  	case "aes-128-cfb":
   114  		plainReader, err = bdcrypto.Aes128CFBDecrypt(bdcrypto.Convert16bytes(key), cipherFile)
   115  	case "aes-192-cfb":
   116  		plainReader, err = bdcrypto.Aes192CFBDecrypt(bdcrypto.Convert24bytes(key), cipherFile)
   117  	case "aes-256-cfb":
   118  		plainReader, err = bdcrypto.Aes256CFBDecrypt(bdcrypto.Convert32bytes(key), cipherFile)
   119  	case "aes-128-ofb":
   120  		plainReader, err = bdcrypto.Aes128OFBDecrypt(bdcrypto.Convert16bytes(key), cipherFile)
   121  	case "aes-192-ofb":
   122  		plainReader, err = bdcrypto.Aes192OFBDecrypt(bdcrypto.Convert24bytes(key), cipherFile)
   123  	case "aes-256-ofb":
   124  		plainReader, err = bdcrypto.Aes256OFBDecrypt(bdcrypto.Convert32bytes(key), cipherFile)
   125  	default:
   126  		return "", fmt.Errorf("unknown decrypt method: %s", method)
   127  	}
   128  
   129  	if err != nil {
   130  		return
   131  	}
   132  
   133  	cipherFileInfo, err := cipherFile.Stat()
   134  	if err != nil {
   135  		return
   136  	}
   137  
   138  	decryptedFilePath = strings.TrimSuffix(filePath, ".encrypt")
   139  	decryptedTmpFilePath := decryptedFilePath + ".decrypted"
   140  	decryptedTmpFile, err := os.OpenFile(decryptedTmpFilePath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, cipherFileInfo.Mode())
   141  	if err != nil {
   142  		return
   143  	}
   144  
   145  	_, err = io.Copy(decryptedTmpFile, plainReader)
   146  	if err != nil {
   147  		return
   148  	}
   149  
   150  	defer decryptedTmpFile.Close()
   151  
   152  	if isGzip {
   153  		err = bdcrypto.GZIPUnompressFile(decryptedTmpFilePath)
   154  		if err != nil {
   155  			os.Remove(decryptedTmpFilePath)
   156  			return
   157  		}
   158  
   159  		// 删除已加密的文件
   160  		os.Remove(filePath)
   161  	}
   162  
   163  	if filePath != decryptedFilePath {
   164  		os.Rename(decryptedTmpFilePath, decryptedFilePath)
   165  	} else {
   166  		decryptedFilePath = decryptedTmpFilePath
   167  	}
   168  
   169  	return decryptedFilePath, nil
   170  }