crypto package updated

crypto/decryption.go

@@ -5,10 +5,36 @@ import (
 	"crypto/cipher"
 	"encoding/base64"
 	"errors"
+	"strconv"
+	"strings"
 )
 
 // Decrypt function is used to decrypt the string
 func Decrypt(securemess string) (decodedmess string, err error) {
+	if len(strings.TrimSpace(securemess)) == 0 {
+		return "", errors.New("string is empty")
+	}
+	decodedStr := strings.Split(securemess, strconv.Itoa(encV1)+"||")
+	if len(decodedStr) == 2 {
+		ver, err := strconv.Atoi(decodedStr[0])
+		if err != nil {
+			return "", err
+		}
+		switch ver {
+		case encV1:
+			decodedmess, err = decrypt1(decodedStr[1])
+			if err != nil {
+				return "", err
+			}
+		default:
+			return "", errors.New("invalid encryption")
+		}
+	}
+
+	return decodedmess, nil
+}
+
+func decrypt1(securemess string) (string, error) {
 	cipherText, err := base64.URLEncoding.DecodeString(securemess)
 	if err != nil {
 		return "", err
@@ -33,6 +59,6 @@ func Decrypt(securemess string) (decodedmess string, err error) {
 	// XORKeyStream can work in-place if the two arguments are the same.
 	stream.XORKeyStream(cipherText, cipherText)
 
-	decodedmess = string(cipherText)
+	decodedmess := string(cipherText)
 	return decodedmess, nil
 }

crypto/decryption_test.go

@@ -1,6 +1,7 @@
 package crypto
 
 import (
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -17,7 +18,8 @@ func TestDecrypt(t *testing.T) {
 		assert.Equal(t, TestString, actual)
 	})
 	t.Run("failure decryption", func(t *testing.T) {
-		_, err := Decrypt(TestString)
+		str := fmt.Sprintf("%d%s%s", encV1, "||", TestString)
+		_, err := Decrypt(str)
 		assert.NotNil(t, err)
 	})
 }

crypto/encryption.go

@@ -5,13 +5,23 @@ import (
 	"crypto/cipher"
 	"crypto/rand"
 	"encoding/base64"
+	"errors"
 	"io"
+	"strconv"
+	"strings"
 )
 
 var cipherKey = []byte("0123456789012345")
 
+const (
+	encV1 = 1
+)
+
 //Encrypt function is used to encrypt the string
 func Encrypt(message string) (encmess string, err error) {
+	if len(strings.TrimSpace(message)) == 0 {
+		return "", errors.New("string is empty")
+	}
 	plainText := []byte(message)
 
 	block, err := aes.NewCipher(cipherKey)
@@ -32,5 +42,6 @@ func Encrypt(message string) (encmess string, err error) {
 
 	//returns to base64 encoded string
 	encmess = base64.URLEncoding.EncodeToString(cipherText)
-	return encmess, nil
+	finalEnc := strconv.Itoa(encV1) + "||" + encmess //fmt.Sprintf("%d%s%s", encV1, "||", encmess)
+	return finalEnc, nil
 }