update

aes.go

@@ -19,7 +19,6 @@
 package codec
 
 import (
-	"bytes"
 	"strings"
 )
 
@@ -53,45 +52,3 @@ func NewAES(keyTypes ...string) *AES {
 type AES struct {
 	Codec
 }
-
-const (
-	aes128KeyLen = 128 / 8 // 16
-	aes192KeyLen = 192 / 8 // 24
-	aes256KeyLen = 256 / 8 // 32
-)
-
-// AESKeyTypes AES Key类型
-var AESKeyTypes = map[string]int{
-	`AES-128`: aes128KeyLen,
-	`AES-192`: aes192KeyLen,
-	`AES-256`: aes256KeyLen,
-}
-
-func GenAESKey(key []byte, typ ...string) []byte {
-	var keyType string
-	if len(typ) > 0 {
-		keyType = typ[0]
-	}
-	keyLen, ok := AESKeyTypes[keyType]
-	if !ok {
-		keyLen = aes128KeyLen
-	}
-	return FixedAESKey(keyLen, key)
-}
-
-var FixedAESKey = FixedKeyByWhitespacePrefix
-
-func FixedKeyByWhitespacePrefix(keyLen int, key []byte) []byte {
-	if len(key) == keyLen {
-		return key
-	}
-	k := make([]byte, keyLen)
-	if len(key) < keyLen {
-		remains := keyLen - len(key)
-		copy(k, bytes.Repeat([]byte(` `), remains))
-		copy(k[remains:], key)
-	} else {
-		copy(k, key)
-	}
-	return k
-}

aes_algo.go

@@ -0,0 +1,10 @@
+package codec
+
+const (
+	AlgoAES128CBC = `AES-128-CBC`
+	AlgoAES192CBC = `AES-192-CBC`
+	AlgoAES256CBC = `AES-256-CBC`
+	AlgoAES128ECB = `AES-128-ECB`
+	AlgoAES192ECB = `AES-192-ECB`
+	AlgoAES256ECB = `AES-256-ECB`
+)

aes_cbc.go

@@ -8,29 +8,18 @@ import (
 )
 
 func NewAESCBC(keyTypes ...string) *AESCBC {
-	var keyType string
-	if len(keyTypes) > 0 {
-		keyType = keyTypes[0]
-	}
-	return &AESCBC{key: NewSafeKeys(), keyType: keyType}
+	return &AESCBC{aesKey: newAESKey(keyTypes...)}
 }
 
 type AESCBC struct {
-	key     *SafeKeys
-	keyType string
+	*aesKey
 }
 
 func (c *AESCBC) genKey(key []byte) []byte {
-	if c.key == nil {
-		c.key = NewSafeKeys()
-	}
-	ckey := string(key)
-	k, ok := c.key.Get(ckey)
-	if !ok {
-		k = GenAESKey(key, c.keyType)
-		c.key.Set(ckey, k)
+	if c.aesKey == nil {
+		c.aesKey = newAESKey()
 	}
-	return k
+	return c.GetKey(key)
 }
 
 func (c *AESCBC) Encode(rawData, authKey string) string {

aes_ecb.go

@@ -7,29 +7,18 @@ import (
 )
 
 func NewAESECB(keyTypes ...string) *AESECB {
-	var keyType string
-	if len(keyTypes) > 0 {
-		keyType = keyTypes[0]
-	}
-	return &AESECB{key: NewSafeKeys(), keyType: keyType}
+	return &AESECB{aesKey: newAESKey(keyTypes...)}
 }
 
 type AESECB struct {
-	key     *SafeKeys
-	keyType string
+	*aesKey
 }
 
 func (c *AESECB) genKey(key []byte) []byte {
-	if c.key == nil {
-		c.key = NewSafeKeys()
-	}
-	ckey := string(key)
-	k, ok := c.key.Get(ckey)
-	if !ok {
-		k = GenAESKey(key, c.keyType)
-		c.key.Set(ckey, k)
+	if c.aesKey == nil {
+		c.aesKey = newAESKey()
 	}
-	return k
+	return c.GetKey(key)
 }
 
 func (c *AESECB) Encode(rawData, authKey string) string {

aes_key.go

@@ -0,0 +1,99 @@
+package codec
+
+import "bytes"
+
+const (
+	KeyAES128 = `AES-128`
+	KeyAES192 = `AES-192`
+	KeyAES256 = `AES-256`
+)
+
+const (
+	aes128KeyLen = 128 / 8 // 16
+	aes192KeyLen = 192 / 8 // 24
+	aes256KeyLen = 256 / 8 // 32
+)
+
+// AESKeyTypes AES Key类型
+var AESKeyTypes = map[string]int{
+	KeyAES128: aes128KeyLen,
+	KeyAES192: aes192KeyLen,
+	KeyAES256: aes256KeyLen,
+}
+
+type KeyFixer func(keyLen int, key []byte) []byte
+
+func newAESKey(keyTypes ...string) *aesKey {
+	var keyType string
+	if len(keyTypes) > 0 {
+		keyType = keyTypes[0]
+	}
+	return &aesKey{key: NewSafeKeys(), keyType: keyType, keyFixer: FixedAESKey}
+}
+
+type aesKey struct {
+	key      *SafeKeys
+	keyType  string
+	keyFixer KeyFixer
+}
+
+func (a *aesKey) SetKeyFixer(fixer KeyFixer) {
+	a.keyFixer = fixer
+}
+
+func (a *aesKey) GetKey(key []byte) []byte {
+	ckey := string(key)
+	k, ok := a.key.Get(ckey)
+	if !ok {
+		k = a.GenKey(key, a.keyType)
+		a.key.Set(ckey, k)
+	}
+	return k
+}
+
+func (a *aesKey) GenKey(key []byte, typ ...string) []byte {
+	var keyType string
+	if len(typ) > 0 {
+		keyType = typ[0]
+	}
+	keyLen, ok := AESKeyTypes[keyType]
+	if !ok {
+		keyLen = aes128KeyLen
+	}
+	if a.keyFixer == nil {
+		return FixedAESKey(keyLen, key)
+	}
+	return a.keyFixer(keyLen, key)
+}
+
+var FixedAESKey = FixedKeyDefault
+
+func FixedKeyDefault(keyLen int, key []byte) []byte {
+	if len(key) == keyLen {
+		return key
+	}
+
+	k := make([]byte, keyLen)
+	copy(k, key)
+	for i := keyLen; i < len(key); {
+		for j := 0; j < keyLen && i < len(key); j, i = j+1, i+1 {
+			k[j] ^= key[i]
+		}
+	}
+	return k
+}
+
+func FixedKeyByWhitespacePrefix(keyLen int, key []byte) []byte {
+	if len(key) == keyLen {
+		return key
+	}
+	k := make([]byte, keyLen)
+	if len(key) < keyLen {
+		remains := keyLen - len(key)
+		copy(k, bytes.Repeat([]byte(` `), remains))
+		copy(k[remains:], key)
+	} else {
+		copy(k, key)
+	}
+	return k
+}

aes_test.go

@@ -1,12 +1,13 @@
 package codec
 
 import (
+	"encoding/hex"
 	"fmt"
 	"testing"
 	"time"
 )
 
-func TestAes(t *testing.T) {
+func TestAES(t *testing.T) {
 	var (
 		paykey = fmt.Sprintf(`%032d`, time.Now().UnixMicro())
 	)
@@ -18,30 +19,48 @@ func TestAes(t *testing.T) {
 		t.Fatal(`failed`)
 	}
 	//t.Fatal(crypto.Decode(crypted, paykey))
+}
+
+func TestFixedKeyDefault(t *testing.T) {
+	pwd := `123`
+	key := Md5bytes([]byte(pwd))
+	key2 := Md5str(pwd)
+	if string(key) != key2 {
+		t.Fatalf(`%s != %s`, key, key2)
+	}
+	t.Logf(`md5key: %s`, key)
+	fixedKey := FixedKeyDefault(16, key)
+	t.Logf(`fixedKey: %s`, hex.EncodeToString(fixedKey))
 
-	GenAESKey([]byte(`2`))
+	pwd = `1234`
+	key = Md5bytes([]byte(pwd))
+	t.Logf(`md5key: %s`, key)
+	fixedKey = FixedKeyDefault(16, key)
+	t.Logf(`fixedKey: %s`, hex.EncodeToString(fixedKey))
 }
 
 func TestGenAESKey(t *testing.T) {
-	key := GenAESKey([]byte(`2`))
+	aesKey := newAESKey(KeyAES128)
+	aesKey.SetKeyFixer(FixedKeyByWhitespacePrefix)
+	key := aesKey.GenKey([]byte(`2`))
 	t.Logf(`key: [%s]`, key)
 	if string(key) != `               2` {
 		t.Fatal(`failed`)
 	}
 
-	key = GenAESKey([]byte(`123`))
+	key = aesKey.GenKey([]byte(`123`))
 	t.Logf(`key: [%s]`, key)
 	if string(key) != `             123` {
 		t.Fatal(`failed`)
 	}
 
-	key = GenAESKey([]byte(``))
+	key = aesKey.GenKey([]byte(``))
 	t.Logf(`key: [%s]`, key)
 	if string(key) != `                ` {
 		t.Fatal(`failed`)
 	}
 
-	key = GenAESKey([]byte(`12345678901234567890`))
+	key = aesKey.GenKey([]byte(`12345678901234567890`))
 	t.Logf(`key: [%s]`, key)
 	if string(key) != `1234567890123456` {
 		t.Fatal(`failed`)

codec.go

@@ -23,6 +23,7 @@ type Codec interface {
 	Decode(cryptedData, authKey string) string
 	EncodeBytes(rawData, authKey []byte) []byte
 	DecodeBytes(cryptedData, authKey []byte) []byte
+	SetKeyFixer(fixer KeyFixer)
 }
 
 var Default Codec = NewAES()

des.go

@@ -18,39 +18,7 @@
 
 package codec
 
-import (
-	"bytes"
-	"crypto/des"
-)
-
 var (
 	_ Codec = NewDESCBC()
 	_ Codec = NewDESECB()
 )
-
-func GenDESKey(key []byte) []byte {
-	return FixedDESKey(des.BlockSize, key)
-}
-
-var FixedDESKey = FixedKeyByRepeatContent
-
-func FixedKeyByRepeatContent(keyLen int, key []byte) []byte {
-	if len(key) == keyLen {
-		return key
-	}
-	k := make([]byte, keyLen)
-	length := len(key)
-	if length == 0 {
-		copy(k, bytes.Repeat([]byte(` `), keyLen))
-	} else if length < keyLen {
-		div := keyLen / length
-		mod := keyLen % length
-		for i := 0; i < div; i++ {
-			copy(k[length*i:], key)
-		}
-		copy(k[length*div:], key[:mod])
-	} else {
-		copy(k, key)
-	}
-	return k
-}

des_cbc.go

@@ -27,12 +27,12 @@ import (
 
 func NewDESCBC() *DESCBC {
 	return &DESCBC{
-		key: NewSafeKeys(),
+		desKey: newDESKey(),
 	}
 }
 
 type DESCBC struct {
-	key *SafeKeys
+	*desKey
 }
 
 // Encode DES CBC加密
@@ -89,14 +89,8 @@ func (d *DESCBC) DecodeBytes(crypted, secret []byte) []byte {
 }
 
 func (d *DESCBC) GenKey(key []byte) []byte {
-	if d.key == nil {
-		d.key = NewSafeKeys()
+	if d.desKey == nil {
+		d.desKey = newDESKey()
 	}
-	ckey := string(key)
-	kkey, ok := d.key.Get(ckey)
-	if !ok {
-		kkey = GenDESKey(key)
-		d.key.Set(ckey, kkey)
-	}
-	return kkey
+	return d.GetKey(key)
 }