package main

// Sealing to one recipient (novox/hq ADR 0183, ADR 0206): the manager seals what it hands a consumer to
// the key that consumer sent, and a node seals a waiting login to the key the manager gives. The same box
// the agent module's TypeScript makes and opens, byte for byte — X25519 for the agreement, HKDF-SHA256 for
// the key, AES-256-GCM for the box — so `testdata/sealed-by-typescript.json` is opened here, and a test
// reopens what this seals with the same derivation.
//
// A box is `{ v: 1, eph, iv, tag, ct }`, every field base64; `eph` is the one-time public key as SPKI DER,
// and the key is bound to it and to the recipient's raw public key, so a box cannot be re-addressed.

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/ecdh"
	"crypto/hkdf"
	"crypto/rand"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"errors"
	"fmt"
)

// SealedBox is a value sealed to one recipient.
type SealedBox struct {
	V   int    `json:"v"`
	Eph string `json:"eph"`
	IV  string `json:"iv"`
	Tag string `json:"tag"`
	Ct  string `json:"ct"`
}

// KeyPair is a recipient's keypair as the two PEM strings it is kept and sent as.
type KeyPair struct {
	PublicKey  string `json:"publicKey"`
	PrivateKey string `json:"privateKey"`
}

const sealInfo = "novox-mesh sealed box v1"

// GenerateKeyPair makes an X25519 keypair, PEM-encoded as the agent module's are.
func GenerateKeyPair() (KeyPair, error) {
	priv, err := ecdh.X25519().GenerateKey(rand.Reader)
	if err != nil {
		return KeyPair{}, err
	}
	pubDER, err := x509.MarshalPKIXPublicKey(priv.PublicKey())
	if err != nil {
		return KeyPair{}, err
	}
	privDER, err := x509.MarshalPKCS8PrivateKey(priv)
	if err != nil {
		return KeyPair{}, err
	}
	return KeyPair{
		PublicKey:  string(pem.EncodeToMemory(&pem.Block{Type: "PUBLIC KEY", Bytes: pubDER})),
		PrivateKey: string(pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: privDER})),
	}, nil
}

func publicFromPEM(p string) (*ecdh.PublicKey, error) {
	block, _ := pem.Decode([]byte(p))
	if block == nil {
		return nil, errors.New("not a PEM public key")
	}
	k, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, err
	}
	pub, ok := k.(*ecdh.PublicKey)
	if !ok || pub.Curve() != ecdh.X25519() {
		return nil, errors.New("not an X25519 public key")
	}
	return pub, nil
}

func privateFromPEM(p string) (*ecdh.PrivateKey, error) {
	block, _ := pem.Decode([]byte(p))
	if block == nil {
		return nil, errors.New("not a PEM private key")
	}
	k, err := x509.ParsePKCS8PrivateKey(block.Bytes)
	if err != nil {
		return nil, err
	}
	priv, ok := k.(*ecdh.PrivateKey)
	if !ok || priv.Curve() != ecdh.X25519() {
		return nil, errors.New("not an X25519 private key")
	}
	return priv, nil
}

func boxKey(secret, ephDER, recipientRaw []byte) ([]byte, error) {
	salt := append(append([]byte{}, ephDER...), recipientRaw...)
	return hkdf.Key(sha256.New, secret, salt, sealInfo, 32)
}

// Seal seals plaintext to the recipient's public key.
func Seal(plaintext, recipientPEM string) (SealedBox, error) {
	recipient, err := publicFromPEM(recipientPEM)
	if err != nil {
		return SealedBox{}, err
	}
	eph, err := ecdh.X25519().GenerateKey(rand.Reader)
	if err != nil {
		return SealedBox{}, err
	}
	secret, err := eph.ECDH(recipient)
	if err != nil {
		return SealedBox{}, err
	}
	ephDER, err := x509.MarshalPKIXPublicKey(eph.PublicKey())
	if err != nil {
		return SealedBox{}, err
	}
	key, err := boxKey(secret, ephDER, recipient.Bytes())
	if err != nil {
		return SealedBox{}, err
	}
	gcm, err := newGCM(key)
	if err != nil {
		return SealedBox{}, err
	}
	iv := make([]byte, 12)
	if _, err := rand.Read(iv); err != nil {
		return SealedBox{}, err
	}
	out := gcm.Seal(nil, iv, []byte(plaintext), nil)
	ct, tag := out[:len(out)-gcm.Overhead()], out[len(out)-gcm.Overhead():]
	b64 := base64.StdEncoding.EncodeToString
	return SealedBox{V: 1, Eph: b64(ephDER), IV: b64(iv), Tag: b64(tag), Ct: b64(ct)}, nil
}

// Open opens a box with the recipient's private key; it fails for a box to another key or one tampered with.
func Open(box SealedBox, privatePEM string) (string, error) {
	if box.V != 1 {
		return "", errors.New("not a sealed box this module can open")
	}
	priv, err := privateFromPEM(privatePEM)
	if err != nil {
		return "", err
	}
	d := base64.StdEncoding.DecodeString
	ephDER, err := d(box.Eph)
	if err != nil {
		return "", fmt.Errorf("the box's eph: %w", err)
	}
	ephKey, err := x509.ParsePKIXPublicKey(ephDER)
	if err != nil {
		return "", err
	}
	eph, ok := ephKey.(*ecdh.PublicKey)
	if !ok {
		return "", errors.New("the box's eph is not an X25519 key")
	}
	secret, err := priv.ECDH(eph)
	if err != nil {
		return "", err
	}
	key, err := boxKey(secret, ephDER, priv.PublicKey().Bytes())
	if err != nil {
		return "", err
	}
	iv, err1 := d(box.IV)
	tag, err2 := d(box.Tag)
	ct, err3 := d(box.Ct)
	if err := errors.Join(err1, err2, err3); err != nil {
		return "", err
	}
	gcm, err := newGCM(key)
	if err != nil {
		return "", err
	}
	plain, err := gcm.Open(nil, iv, append(ct, tag...), nil)
	if err != nil {
		return "", errors.New("the box does not open with this key")
	}
	return string(plain), nil
}

func newGCM(key []byte) (cipher.AEAD, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	return cipher.NewGCM(block)
}