package main

import (
	"crypto/aes"
	"crypto/cipher"
	"encoding/base64"
	"log"
	"net/http"

	"github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}

var aesKey = []byte("0123456789012345") // 16 bytes AES key

func main() {
	http.HandleFunc("/", handleWebSocket)
	log.Fatal(http.ListenAndServe(":8080", nil))
}

func handleWebSocket(w http.ResponseWriter, r *http.Request) {
	// Upgrade the HTTP connection to a WebSocket connection
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}
	defer conn.Close()

	// Generate a new AES cipher block using the key
	block, err := aes.NewCipher(aesKey)
	if err != nil {
		log.Println(err)
		return
	}

	// Generate a new GCM AEAD cipher using the block
	aead, err := cipher.NewGCM(block)
	if err != nil {
		log.Println(err)
		return
	}

	// Read and write messages on the WebSocket connection
	for {
		// Read a message from the WebSocket connection
		_, msg, err := conn.ReadMessage()
		if err != nil {
			log.Println(err)
			break
		}

		// Decrypt the message using the AEAD cipher
		nonceSize := aead.NonceSize()
		if len(msg) < nonceSize {
			log.Println("Invalid message length")
			break
		}
		nonce, ciphertext := msg[:nonceSize], msg[nonceSize:]
		 plaintext, err := aead.Open(nil, nonce, ciphertext, nil)
		if err != nil {
			log.Println(err)
			break
		}
		log.Printf("Received message: %s", plaintext)

		// Encrypt a message using the AEAD cipher
		nonce = make([]byte, nonceSize)
		if _, err := rand.Read(nonce); err != nil {
			log.Println(err)
			break
		}
		ciphertext = aead.Seal(nonce, nonce, plaintext, nil)
		msg = append(nonce, ciphertext...)
		log.Printf("Sending message: %s", msg)

		// Write the encrypted message to the WebSocket connection
		err = conn.WriteMessage(websocket.TextMessage, msg)
		if err != nil {
			log.Println(err)
			break
		}
	}
}