From dadcad2ef7fcc2d4c914bc88959b69837ef7f5e9 Mon Sep 17 00:00:00 2001
From: Wesley Yan Soares Brehmer
 <102262759+NervousGroove@users.noreply.github.com>
Date: Mon, 11 Sep 2023 16:02:43 -0300
Subject: [PATCH] Update w7.go

---
 w7.go | 75 ++++++++++++++++++++++++++++++++++++++++++++++-------------
 1 file changed, 59 insertions(+), 16 deletions(-)

diff --git a/w7.go b/w7.go
index 9d1253f..12d7148 100644
--- a/w7.go
+++ b/w7.go
@@ -1,35 +1,78 @@
 package w7dth
 
 import (
-	"crypto/sha256"
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/rand"
 	"encoding/hex"
-	"math/rand"
-	"time"
+	"errors"
+	"io"
 )
 
-type W7DTH struct{}
+type W7DTH struct {
+	key []byte
+}
 
 func New() *W7DTH {
 	return &W7DTH{}
 }
 
-func (w *W7DTH) key(generate string) string {
+func (w *W7DTH) Key(generate string) (string, error) {
 	if generate == "generate" {
-		rand.Seed(time.Now().UnixNano())
 		key := make([]byte, 32)
-		rand.Read(key)
-		return hex.EncodeToString(key)
+		_, err := rand.Read(key)
+		if err != nil {
+			return "", err
+		}
+		w.key = key
+		return hex.EncodeToString(key), nil
+	}
+	key, err := hex.DecodeString(generate)
+	if err != nil {
+		return "", err
 	}
-	return generate
+	w.key = key
+	return generate, nil
 }
 
-func (w *W7DTH) Encrypt(data string, key string) string {
-	hash := sha256.New()
-	hash.Write([]byte(data + key))
-	return hex.EncodeToString(hash.Sum(nil))
+func (w *W7DTH) Encrypt(data string) (string, error) {
+	block, err := aes.NewCipher(w.key)
+	if err != nil {
+		return "", err
+	}
+	gcm, err := cipher.NewGCM(block)
+	if err != nil {
+		return "", err
+	}
+	nonce := make([]byte, gcm.NonceSize())
+	if _, err = io.ReadFull(rand.Reader, nonce); err != nil {
+		return "", err
+	}
+	ciphertext := gcm.Seal(nonce, nonce, []byte(data), nil)
+	return hex.EncodeToString(ciphertext), nil
 }
 
-func (w *W7DTH) Decrypt(data string, key string) string {
-	// SHA-256 is a one-way function, it cannot be decrypted
-	return ""
+func (w *W7DTH) Decrypt(data string) (string, error) {
+	ciphertext, err := hex.DecodeString(data)
+	if err != nil {
+		return "", err
+	}
+	block, err := aes.NewCipher(w.key)
+	if err != nil {
+		return "", err
+	}
+	gcm, err := cipher.NewGCM(block)
+	if err != nil {
+		return "", err
+	}
+	nonceSize := gcm.NonceSize()
+	if len(ciphertext) < nonceSize {
+		return "", errors.New("ciphertext too short")
+	}
+	nonce, ciphertext := ciphertext[:nonceSize], ciphertext[nonceSize:]
+	plaintext, err := gcm.Open(nil, nonce, ciphertext, nil)
+	if err != nil {
+		return "", err
+	}
+	return string(plaintext), nil
 }