/
aes.go
136 lines (120 loc) · 3.52 KB
/
aes.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
package cryptogy
import (
"bytes"
"crypto/aes"
"crypto/cipher"
"strings"
)
type (
PaddingFunc func(cipherText []byte, blockSize int) []byte
UnpaddingFunc func(origData []byte) []byte
)
func ZeroPadding(cipherText []byte, blockSize int) []byte {
padding := blockSize - len(cipherText)%blockSize
padtext := bytes.Repeat([]byte{0}, padding) // 用0去填充
return append(cipherText, padtext...)
}
func ZeroUnpadding(origData []byte) []byte {
return bytes.TrimFunc(origData, func(r rune) bool {
return r == rune(0)
})
}
func PKCS5Padding(cipherText []byte, blockSize int) []byte {
padding := blockSize - len(cipherText)%blockSize // 需要padding的数目
// 只要少于256就能放到一个byte中,默认的blockSize=16(即采用16*8=128, AES-128长的密钥)
// 最少填充1个byte,如果原文刚好是blocksize的整数倍,则再填充一个blocksize
padtext := bytes.Repeat([]byte{byte(padding)}, padding) // 生成填充的文本
return append(cipherText, padtext...)
}
func PKCS5Unpadding(origData []byte) []byte {
length := len(origData)
unpadding := int(origData[length-1])
return origData[:(length - unpadding)]
}
func PKCS7Padding(cipherText []byte, blockSize int) []byte {
padding := blockSize - len(cipherText)%blockSize
padtext := bytes.Repeat([]byte{byte(padding)}, padding)
return append(cipherText, padtext...)
}
func PKCS7Unpadding(origData []byte) []byte {
length := len(origData)
unpadding := int(origData[length-1])
return origData[:(length - unpadding)]
}
// AES加密,支持模式CBC、CFB、CTR、OFB,不支持ECB和GCM
// 其中CBC模式一般需要填充,用法: c.SetPaddingFunc("PKCS5")
type AESCipher struct {
modeName string
iv []byte
Padding PaddingFunc
Unpadding UnpaddingFunc
cipher.Block
}
func NewAESCipher(mode string, key []byte) (*AESCipher, error) {
c := &AESCipher{modeName: strings.ToUpper(mode)}
block, err := aes.NewCipher(key)
if err == nil {
c.Block = block
c.iv = key[:c.BlockSize()]
}
return c, err
}
func (c *AESCipher) SetPaddingFunc(name string) {
switch strings.ToUpper(name) {
case "0", "ZERO":
c.Padding = ZeroPadding
c.Unpadding = ZeroUnpadding
case "5", "PKCS5":
c.Padding = PKCS5Padding
c.Unpadding = PKCS5Unpadding
case "7", "PKCS7":
c.Padding = PKCS7Padding
c.Unpadding = PKCS7Unpadding
}
}
func (c *AESCipher) GetStream(isDecrypt bool) cipher.Stream {
switch c.modeName {
default:
return nil
case "CFB":
if isDecrypt {
return cipher.NewCFBDecrypter(c.Block, c.iv)
} else {
return cipher.NewCFBEncrypter(c.Block, c.iv)
}
case "CTR":
return cipher.NewCTR(c.Block, c.iv)
case "OFB":
return cipher.NewOFB(c.Block, c.iv)
}
}
func (c *AESCipher) GetEncrypter() cipher.BlockMode {
return cipher.NewCBCEncrypter(c.Block, c.iv)
}
func (c *AESCipher) GetDecrypter() cipher.BlockMode {
return cipher.NewCBCDecrypter(c.Block, c.iv)
}
func (c *AESCipher) Encrypt(origData []byte) ([]byte, error) {
if c.Padding != nil {
origData = c.Padding(origData, c.BlockSize())
}
cipherText := make([]byte, len(origData))
if c.modeName == "CBC" {
c.GetEncrypter().CryptBlocks(cipherText, origData)
} else {
c.GetStream(false).XORKeyStream(cipherText, origData)
}
return cipherText, nil
}
func (c *AESCipher) Decrypt(cipherText []byte) ([]byte, error) {
origData := make([]byte, len(cipherText))
if c.modeName == "CBC" {
c.GetDecrypter().CryptBlocks(origData, cipherText)
} else {
c.GetStream(true).XORKeyStream(cipherText, origData)
}
if c.Unpadding != nil {
origData = c.Unpadding(origData)
}
return origData, nil
}