/
encrypt_stream.go
95 lines (83 loc) · 2.33 KB
/
encrypt_stream.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
package encrypt
import (
"crypto/cipher"
)
// StreamOption option
type StreamOption func(bs *blockStream)
func WithStreamRandIV(generateIv func(block cipher.Block) ([]byte, error)) StreamOption {
return func(bs *blockStream) {
bs.generateIv = generateIv
}
}
func WithStreamCodec(newEncrypt, newDecrypt func(block cipher.Block, iv []byte) cipher.Stream) StreamOption {
return func(bs *blockStream) {
bs.newEncrypt = newEncrypt
bs.newDecrypt = newDecrypt
}
}
// NewStreamCipher new with newCipher and key
// newCipher support follow or implement func(key []byte) (cipher.Block, error):
// aes
// cipher
// des
// blowfish
// cast5
// twofish
// xtea
// tea
// block stream cipher
// support:
// cfb(default): cipher.NewCFBEncrypter, cipher.NewCFBDecrypter
// ctr: cipher.NewCTR, cipher.NewCTR
// ofb: cipher.NewOFB, cipher.NewOFB
func NewStreamCipher(key []byte, newCipher func(key []byte) (cipher.Block, error), opts ...StreamOption) (BlockCrypt, error) {
block, err := newCipher(key)
if err != nil {
return nil, err
}
bs := &blockStream{
block: block,
newEncrypt: cipher.NewCFBEncrypter,
newDecrypt: cipher.NewCFBDecrypter,
}
for _, opt := range opts {
opt(bs)
}
return bs, nil
}
type blockStream struct {
block cipher.Block
generateIv func(block cipher.Block) ([]byte, error)
newEncrypt func(block cipher.Block, iv []byte) cipher.Stream
newDecrypt func(block cipher.Block, iv []byte) cipher.Stream
}
func (sf *blockStream) BlockSize() int {
return sf.block.BlockSize()
}
func (sf *blockStream) Encrypt(plainText []byte) ([]byte, error) {
if len(plainText) == 0 {
return nil, ErrInputInvalidLength
}
blockSize := sf.block.BlockSize()
ivFunc := RandIV
if sf.generateIv != nil {
ivFunc = sf.generateIv
}
iv, err := ivFunc(sf.block)
if err != nil || len(iv) != blockSize {
return nil, ErrInvalidIvSize
}
cipherText := make([]byte, blockSize+len(plainText))
copy(cipherText[:blockSize], iv)
sf.newEncrypt(sf.block, iv).XORKeyStream(cipherText[blockSize:], plainText)
return cipherText, nil
}
func (sf *blockStream) Decrypt(cipherText []byte) ([]byte, error) {
blockSize := sf.block.BlockSize()
if len(cipherText) < blockSize {
return nil, ErrInputNotMoreABlock
}
iv, msg := cipherText[:blockSize], cipherText[blockSize:]
sf.newDecrypt(sf.block, iv).XORKeyStream(msg, msg)
return msg, nil
}