New crypto.EncryptAES/DecryptAES functions (#1305)

Signed-off-by: Dave Henderson <dhenderson@gmail.com>

2 files changed, 81 insertions, 0 deletions

diff --git a/funcs/crypto.go b/funcs/crypto.go
index 99d2cce5..8fc28f69 100644
--- a/funcs/crypto.go
+++ b/funcs/crypto.go
@@ -197,3 +197,50 @@ func (f *CryptoFuncs) RSADerivePublicKey(privateKey string) (string, error) {
 	out, err := crypto.RSADerivePublicKey([]byte(privateKey))
 	return string(out), err
 }
+
+// EncryptAES -
+func (f *CryptoFuncs) EncryptAES(key string, args ...interface{}) ([]byte, error) {
+	k, msg, err := parseAESArgs(key, args...)
+	if err != nil {
+		return nil, err
+	}
+
+	return crypto.EncryptAESCBC(k, msg)
+}
+
+// DecryptAES -
+func (f *CryptoFuncs) DecryptAES(key string, args ...interface{}) (string, error) {
+	out, err := f.DecryptAESBytes(key, args...)
+	return conv.ToString(out), err
+}
+
+// DecryptAESBytes -
+func (f *CryptoFuncs) DecryptAESBytes(key string, args ...interface{}) ([]byte, error) {
+	k, msg, err := parseAESArgs(key, args...)
+	if err != nil {
+		return nil, err
+	}
+
+	return crypto.DecryptAESCBC(k, msg)
+}
+
+func parseAESArgs(key string, args ...interface{}) ([]byte, []byte, error) {
+	keyBits := 256 // default to AES-256-CBC
+
+	var msg []byte
+
+	switch len(args) {
+	case 1:
+		msg = toBytes(args[0])
+	case 2:
+		keyBits = conv.ToInt(args[0])
+		msg = toBytes(args[1])
+	default:
+		return nil, nil, fmt.Errorf("wrong number of args: want 2 or 3, got %d", len(args))
+	}
+
+	k := make([]byte, keyBits/8)
+	copy(k, []byte(key))
+
+	return k, msg, nil
+}
diff --git a/funcs/crypto_test.go b/funcs/crypto_test.go
index 4cb968c0..5faa3c6a 100644
--- a/funcs/crypto_test.go
+++ b/funcs/crypto_test.go
@@ -138,3 +138,37 @@ func TestRSACrypt(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Equal(t, dec, string(b))
 }
+
+func TestAESCrypt(t *testing.T) {
+	c := testCryptoNS()
+	key := "0123456789012345"
+	in := "hello world"
+
+	_, err := c.EncryptAES(key, 1, 2, 3, 4)
+	assert.Error(t, err)
+
+	_, err = c.DecryptAES(key, 1, 2, 3, 4)
+	assert.Error(t, err)
+
+	enc, err := c.EncryptAES(key, in)
+	assert.NoError(t, err)
+
+	dec, err := c.DecryptAES(key, enc)
+	assert.NoError(t, err)
+	assert.Equal(t, in, dec)
+
+	b, err := c.DecryptAESBytes(key, enc)
+	assert.NoError(t, err)
+	assert.Equal(t, dec, string(b))
+
+	enc, err = c.EncryptAES(key, 128, in)
+	assert.NoError(t, err)
+
+	dec, err = c.DecryptAES(key, 128, enc)
+	assert.NoError(t, err)
+	assert.Equal(t, in, dec)
+
+	b, err = c.DecryptAESBytes(key, 128, enc)
+	assert.NoError(t, err)
+	assert.Equal(t, dec, string(b))
+}