From 92a2259f302763009a7b8e15711aa71aeee8d6ee Mon Sep 17 00:00:00 2001
From: darmiel <71837281+darmiel@users.noreply.github.com>
Date: Thu, 25 Mar 2021 14:11:01 +0100
Subject: [PATCH] Implemented encryption (#2)

---
 internal/common/encryption.go | 86 +++++++++++++++++++++++++++++++++++
 internal/server/request.go    | 33 +++++++++++---
 main.go                       |  4 +-
 3 files changed, 116 insertions(+), 7 deletions(-)
 create mode 100644 internal/common/encryption.go

diff --git a/internal/common/encryption.go b/internal/common/encryption.go
new file mode 100644
index 0000000..23f8fab
--- /dev/null
+++ b/internal/common/encryption.go
@@ -0,0 +1,86 @@
+package common
+
+import (
+	"bytes"
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/rand"
+	"encoding/base64"
+	"errors"
+	"io"
+	"math"
+)
+
+var (
+	errKeyEmpty    = errors.New("key empty")
+	errKeyLength   = errors.New("invalid key length")
+	errNonceLength = errors.New("invalid nonce length")
+	bytesEmpty     []byte
+)
+
+func Encrypt(textStr, keyStr string) (res []byte, err error) {
+	gcm, err := getGCM(keyStr)
+	if err != nil {
+		return bytesEmpty, err
+	}
+	nonce := make([]byte, gcm.NonceSize())
+	if _, err = io.ReadFull(rand.Reader, nonce); err != nil {
+		return bytesEmpty, err
+	}
+	seal := gcm.Seal(nonce, nonce, []byte(textStr), nil)
+	res = []byte(base64.StdEncoding.EncodeToString(seal))
+	return
+}
+
+func Decrypt(textStr, keyStr string) (res []byte, err error) {
+	text, err := base64.StdEncoding.DecodeString(textStr)
+	if err != nil {
+		return bytesEmpty, err
+	}
+	gcm, err := getGCM(keyStr)
+	if err != nil {
+		return bytesEmpty, err
+	}
+	nonceSize := gcm.NonceSize()
+	if len(text) < nonceSize {
+		return bytesEmpty, errNonceLength
+	}
+	nonce, encrypted := text[:nonceSize], text[nonceSize:]
+	plain, err := gcm.Open(nil, nonce, encrypted, nil)
+	if err != nil {
+		return bytesEmpty, err
+	}
+	return plain, nil
+}
+
+////
+
+func getGCM(keyStr string) (cipher.AEAD, error) {
+	key, err := getKey(keyStr)
+	if err != nil {
+		return nil, err
+	}
+	block, err := aes.NewCipher(key)
+	if err != nil {
+		return nil, err
+	}
+	return cipher.NewGCM(block)
+}
+
+func getKey(keyStr string) (key []byte, err error) {
+	key = []byte(keyStr)
+	if len(key) <= 0 {
+		return key, errKeyEmpty
+	}
+	if len(key) < 32 {
+		rep := math.Ceil(32.0 / float64(len(keyStr)))
+		key = bytes.Repeat([]byte(keyStr), int(rep))
+	}
+	if len(key) > 32 {
+		key = key[:32]
+	}
+	if len(key) != 32 {
+		return key, errKeyLength
+	}
+	return
+}
diff --git a/internal/server/request.go b/internal/server/request.go
index 3a35f61..0ea6471 100644
--- a/internal/server/request.go
+++ b/internal/server/request.go
@@ -1,6 +1,7 @@
 package server
 
 import (
+	"github.com/darmiel/yaxc/internal/common"
 	"github.com/gofiber/fiber/v2"
 	"time"
 )
@@ -8,32 +9,43 @@ import (
 func (s *yAxCServer) handlePostAnywhere(ctx *fiber.Ctx) (err error) {
 	path := ctx.Params("anywhere")
 
+	// Read content
 	bytes := ctx.Body()
 	if s.MaxBodyLength > 0 && len(bytes) > s.MaxBodyLength {
 		return s.errBodyLen
 	}
 	content := string(bytes)
 
-	// custom ttl
+	// TTL
 	ttl := s.DefaultTTL
-
 	if q := ctx.Query("ttl"); q != "" {
 		if ttl, err = time.ParseDuration(q); err != nil {
 			return
 		}
 	}
 
-	// check if ttl is valid
+	// Encryption
+	if q := ctx.Query("secret"); q != "" {
+		// fail on error
+		encrypt, err := common.Encrypt(content, q)
+		if err != nil {
+			return err
+		}
+		content = string(encrypt)
+	}
+
+	// Check if ttl is valid
 	if (s.MinTTL != 0 && s.MinTTL > ttl) || (s.MaxTTL != 0 && s.MaxTTL < ttl) {
 		return ctx.Status(400).SendString("ERROR: TTL out of range")
 	}
 
-	// set contents
+	// Set contents
 	if err := s.Backend.Set(path, content, ttl); err != nil {
 		return ctx.Status(400).SendString("ERROR: " + err.Error())
 	}
 
-	log.Info("Received", content, "to", path, "with a ttl of", ttl)
+	log.Debug(ctx.IP(), "updated", path)
+
 	return ctx.Status(200).SendString(content)
 }
 
@@ -43,7 +55,16 @@ func (s *yAxCServer) handleGetAnywhere(ctx *fiber.Ctx) (err error) {
 	if res, err = s.Backend.Get(path); err != nil {
 		return
 	}
-	log.Info("Send", res, "to", ctx.IP(), "requesting", path)
+
+	// Encryption
+	if q := ctx.Query("secret"); q != "" {
+		// do not fail on error
+		if encrypt, err := common.Decrypt(res, q); err == nil {
+			res = string(encrypt)
+		}
+	}
+
+	log.Debug(ctx.IP(), "requested", path)
 
 	if res == "" {
 		ctx.Status(404)
diff --git a/main.go b/main.go
index a85980c..a783ac5 100644
--- a/main.go
+++ b/main.go
@@ -15,7 +15,9 @@ limitations under the License.
 */
 package main
 
-import "github.com/darmiel/yaxc/cmd"
+import (
+	"github.com/darmiel/yaxc/cmd"
+)
 
 func main() {
 	cmd.Execute()