package main

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"html/template"
	"log"
	"math/big"
	"net/http"
	"time"
)

var (
	priv    *rsa.PrivateKey
	pubPEM  []byte
	certPEM []byte // self-signed cert jen pro sdílení identity (volitelné)
	tmpl    *template.Template
)

type envelope struct {
	// Encrypted AES key, Nonce, Ciphertext (GCM)
	EK string `json:"ek"` // base64(RSA-OAEP(aesKey))
	N  string `json:"n"`  // base64(nonce 12B)
	CT string `json:"ct"` // base64(GCM(ciphertext||tag))
}

func main() {
	var err error
	// 1) identita: RSA klíče
	priv, err = rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		log.Fatal(err)
	}
	pubASN1, _ := x509.MarshalPKIXPublicKey(&priv.PublicKey)
	pubPEM = pem.EncodeToMemory(&pem.Block{Type: "PUBLIC KEY", Bytes: pubASN1})

	// 2) volitelně vygeneruj self-signed cert pro export (není pro HTTPS)
	certPEM = generateSelfSignedCert(&priv.PublicKey)

	// 3) šablony
	tmpl = template.Must(template.ParseGlob("templates/*.html"))

	// 4) routing
	http.HandleFunc("/", indexHandler)
	http.HandleFunc("/public.pem", publicKeyHandler)
	http.HandleFunc("/public.crt", publicCertHandler)
	http.HandleFunc("/encrypt", encryptHandler) // POST
	http.HandleFunc("/decrypt", decryptHandler) // POST
	http.Handle("/static/", http.StripPrefix("/static/", http.FileServer(http.Dir("static"))))

	log.Println("Server běží na http://localhost:8080 (TLS neřešíme; klíč slouží jen k šifrování zpráv).")
	log.Fatal(http.ListenAndServe(":8080", nil))
}

func indexHandler(w http.ResponseWriter, r *http.Request) {
	_ = tmpl.ExecuteTemplate(w, "index.html", nil)
}

func publicKeyHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/x-pem-file")
	w.Write(pubPEM)
}

func publicCertHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/x-pem-file")
	w.Write(certPEM)
}

func encryptHandler(w http.ResponseWriter, r *http.Request) {
	if err := r.ParseForm(); err != nil {
		http.Error(w, "Bad form", http.StatusBadRequest)
		return
	}
	msg := r.Form.Get("message")
	peer := r.Form.Get("pubkey") // může to být PEM PUBLIC KEY nebo CERT

	pubKey, err := parsePeerPublicKey(peer)
	if err != nil {
		http.Error(w, "Neplatný public key/cert: "+err.Error(), http.StatusBadRequest)
		return
	}

	// --- hybrid: AES-GCM + RSA-OAEP(SHA-256) ---
	// 1) vygeneruj náhodný sym. klíč
	aesKey := make([]byte, 32) // AES-256
	if _, err := rand.Read(aesKey); err != nil {
		http.Error(w, "Rand fail", 500)
		return
	}

	block, err := aes.NewCipher(aesKey)
	if err != nil {
		http.Error(w, "AES fail", 500)
		return
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		http.Error(w, "GCM fail", 500)
		return
	}

	nonce := make([]byte, gcm.NonceSize())
	if _, err := rand.Read(nonce); err != nil {
		http.Error(w, "Nonce fail", 500)
		return
	}

	ciphertext := gcm.Seal(nil, nonce, []byte(msg), nil)

	// 2) zašifruj AES klíč cizím RSA klíčem (OAEP SHA-256)
	label := []byte{} // prázdný label
	ek, err := rsa.EncryptOAEP(sha256.New(), rand.Reader, pubKey, aesKey, label)
	if err != nil {
		http.Error(w, "RSA-OAEP fail: "+err.Error(), 500)
		return
	}

	env := envelope{
		EK: base64.StdEncoding.EncodeToString(ek),
		N:  base64.StdEncoding.EncodeToString(nonce),
		CT: base64.StdEncoding.EncodeToString(ciphertext),
	}

	payload, _ := json.MarshalIndent(env, "", "  ")
	_ = tmpl.ExecuteTemplate(w, "encrypt.html", string(payload))
}

func decryptHandler(w http.ResponseWriter, r *http.Request) {
	if err := r.ParseForm(); err != nil {
		http.Error(w, "Bad form", http.StatusBadRequest)
		return
	}
	in := r.Form.Get("payload")

	var env envelope
	if err := json.Unmarshal([]byte(in), &env); err != nil {
		http.Error(w, "Neplatné JSON", 400)
		return
	}

	ek, err := base64.StdEncoding.DecodeString(env.EK)
	if err != nil {
		http.Error(w, "ek base64", 400)
		return
	}
	nonce, err := base64.StdEncoding.DecodeString(env.N)
	if err != nil {
		http.Error(w, "nonce base64", 400)
		return
	}
	ct, err := base64.StdEncoding.DecodeString(env.CT)
	if err != nil {
		http.Error(w, "ct base64", 400)
		return
	}

	// 1) rozšifruj AES klíč naším RSA
	label := []byte{}
	aesKey, err := rsa.DecryptOAEP(sha256.New(), rand.Reader, priv, ek, label)
	if err != nil {
		http.Error(w, "RSA-OAEP decrypt fail: "+err.Error(), 400)
		return
	}

	block, err := aes.NewCipher(aesKey)
	if err != nil {
		http.Error(w, "AES fail", 500)
		return
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		http.Error(w, "GCM fail", 500)
		return
	}
	plain, err := gcm.Open(nil, nonce, ct, nil)
	if err != nil {
		http.Error(w, "GCM open fail: "+err.Error(), 400)
		return
	}

	_ = tmpl.ExecuteTemplate(w, "decrypt.html", string(plain))
}

// parsePeerPublicKey umí vzít buď PEM PUBLIC KEY, nebo PEM CERT a vrátí *rsa.PublicKey
func parsePeerPublicKey(pemOrCert string) (*rsa.PublicKey, error) {
	block, _ := pem.Decode([]byte(pemOrCert))
	if block == nil {
		return nil, errf("nenašel jsem PEM blok")
	}
	switch block.Type {
	case "PUBLIC KEY":
		k, err := x509.ParsePKIXPublicKey(block.Bytes)
		if err != nil {
			return nil, err
		}
		rsaPub, ok := k.(*rsa.PublicKey)
		if !ok {
			return nil, errf("očekávám RSA PUBLIC KEY")
		}
		return rsaPub, nil
	case "CERTIFICATE":
		c, err := x509.ParseCertificate(block.Bytes)
		if err != nil {
			return nil, err
		}
		rsaPub, ok := c.PublicKey.(*rsa.PublicKey)
		if !ok {
			return nil, errf("cert neobsahuje RSA klíč")
		}
		return rsaPub, nil
	default:
		return nil, errf("nepodporovaný PEM typ: %s", block.Type)
	}
}

func generateSelfSignedCert(pub *rsa.PublicKey) []byte {
	tpl := x509.Certificate{
		SerialNumber: big.NewInt(1),
		Subject: pkix.Name{
			CommonName: "Encryptor Local Identity",
		},
		NotBefore: time.Now(),
		NotAfter:  time.Now().AddDate(1, 0, 0),
		KeyUsage:  x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment,
		// Tenhle cert je jen „vizitka“ (není pro TLS).
		BasicConstraintsValid: true,
	}
	der, _ := x509.CreateCertificate(rand.Reader, &tpl, &tpl, pub, priv)
	buf := &bytes.Buffer{}
	_ = pem.Encode(buf, &pem.Block{Type: "CERTIFICATE", Bytes: der})
	return buf.Bytes()
}

// malá helper chyba
type strErr string

func (e strErr) Error() string      { return string(e) }
func errf(s string, a ...any) error { return strErr(fmtS(s, a...)) }
func fmtS(format string, a ...any) string {
	var b bytes.Buffer
	b.WriteString(format)
	if len(a) > 0 {
		b.WriteString(": ")
	}
	for i, v := range a {
		if i > 0 {
			b.WriteString(", ")
		}
		b.WriteString(anyToString(v))
	}
	return b.String()
}
func anyToString(v any) string {
	switch t := v.(type) {
	case string:
		return t
	default:
		j, _ := json.Marshal(t)
		return string(j)
	}
}