package main

import (
	"bufio"
	"crypto/aes"
	"crypto/cipher"
	"encoding/base64"
	"strings"
	"fmt"

	"log"
	"net"
)

func main() {
	host := "0.0.0.0"
	port := 8080

	listener, err := net.Listen("tcp", fmt.Sprintf("%s:%d", host, port))
	if err != nil {
		log.Fatalf("Error listening: %v", err)
	}
	defer listener.Close()

	fmt.Printf("Server listening on %s:%d\n", host, port)

	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Printf("Error accepting connection: %v", err)
			continue
		}

		go handleConnection(conn)
	}
}

func handleConnection(conn net.Conn) {
    defer conn.Close()

    fmt.Println("Got conn")

    keyData, err := bufio.NewReader(conn).ReadString('\n')
    if err != nil {
        log.Printf("Error reading key: %v", err)
        return
    }
    key, _ := base64.StdEncoding.DecodeString(strings.TrimSpace(keyData))
    conn.Write([]byte("Received key\n"))

    ivData, _ := bufio.NewReader(conn).ReadString('\n')
    iv, _ := base64.StdEncoding.DecodeString(strings.TrimSpace(ivData))
    conn.Write([]byte("Received IV\n"))

    var ciphertext []byte
    var chunk string
    var plaintext []byte

    for {
        chunk, err = bufio.NewReader(conn).ReadString('\n')
        if err != nil {
            log.Printf("Error reading chunk: %v", err)
            return
        }
        if strings.TrimSpace(chunk) == "END_OF_COMMUNICATION" {
            fmt.Println("Client ended communication")
            break
        }
        ciphertext, _ = base64.StdEncoding.DecodeString(strings.TrimSpace(chunk))

        plaintextChunk, err := decrypt(ciphertext, key, iv)
        if err != nil {
            log.Printf("Error decrypting chunk: %v", err)
            return
        }
        plaintext = append(plaintext, plaintextChunk...)

        conn.Write([]byte("Received and decrypted chunk\n"))
    }

    fmt.Println("Decrypted text:", string(plaintext))
    conn.Write([]byte("Ready for next operation\n"))
}




func decrypt(cipherText []byte, key []byte, iv []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	if len(cipherText) < aes.BlockSize {
		return nil, fmt.Errorf("ciphertext too short")
	}

	mode := cipher.NewCBCDecrypter(block, iv)
	mode.CryptBlocks(cipherText, cipherText)

	cipherText = PKCS5Unpadding(cipherText)

	return cipherText, nil
}

func PKCS5Unpadding(data []byte) []byte {
	length := len(data)
	unpadding := int(data[length-1])
	return data[:(length - unpadding)]
}