package mls

// MLS group crypto operations (RFC 9420 §11, §12).
// GroupInfo signing, Welcome encryption/decryption.

import "errors"

var (
	errGroupSecretsNotFound = errors.New("mls: group secrets not found for key package ref")
)

func (info *groupInfo) verifySignature(signerPub signaturePublicKey) bool {
	cs := info.groupContext.cipherSuite
	var w Writer
	info.marshalTBS(&w)
	tbs, err := w.bytes()
	if err != nil {
		return false
	}
	return cs.verifyWithLabel(signerPub, []byte("GroupInfoTBS"), tbs, info.signature)
}

func (info *groupInfo) sign(signerPriv signaturePrivateKey) error {
	cs := info.groupContext.cipherSuite
	var w Writer
	info.marshalTBS(&w)
	tbs, err := w.bytes()
	if err != nil {
		return err
	}
	sig, err := cs.signWithLabel(signerPriv, []byte("GroupInfoTBS"), tbs)
	if err != nil {
		return err
	}
	info.signature = sig
	return nil
}

func (info *groupInfo) verifyConfirmationTag(joinerSecret, pskSecret []byte) bool {
	cs := info.groupContext.cipherSuite
	epochSecret, err := info.groupContext.extractEpochSecret(joinerSecret, pskSecret)
	if err != nil {
		return false
	}
	confirmationKey, err := cs.deriveSecret(epochSecret, secretLabelConfirm)
	if err != nil {
		return false
	}
	return cs.verifyMAC(confirmationKey, info.groupContext.confirmedTranscriptHash, info.confirmationTag)
}

func (info *groupInfo) encrypt(joinerSecret, pskSecret []byte) ([]byte, error) {
	cs := info.groupContext.cipherSuite
	welcomeSecret, err := extractWelcomeSecret(cs, joinerSecret, pskSecret)
	if err != nil {
		return nil, err
	}

	welcomeNonce, err := cs.expandWithLabel(welcomeSecret, []byte("nonce"), nil, uint16(cs.AEADNonceSize()))
	if err != nil {
		return nil, err
	}
	welcomeKey, err := cs.expandWithLabel(welcomeSecret, []byte("key"), nil, uint16(cs.AEADKeySize()))
	if err != nil {
		return nil, err
	}

	rawGroupInfo, err := marshalRaw(info)
	if err != nil {
		return nil, err
	}

	return cs.aeadSeal(welcomeKey, welcomeNonce, rawGroupInfo, nil)
}

func (sec *groupSecrets) encrypt(cs CipherSuite, initKey hpkePublicKey, encryptedGroupInfo []byte) (*hpkeCiphertext, error) {
	raw, err := marshalRaw(sec)
	if err != nil {
		return nil, err
	}
	kemOutput, ciphertext, err := cs.encryptWithLabel(initKey, []byte("Welcome"), encryptedGroupInfo, raw)
	if err != nil {
		return nil, err
	}
	return &hpkeCiphertext{kemOutput: kemOutput, ciphertext: ciphertext}, nil
}

func (w *Welcome) decryptGroupSecrets(ref KeyPackageRef, initKeyPriv hpkePrivateKey) (*groupSecrets, error) {
	cs := w.cipherSuite
	sec := w.findSecret(ref)
	if sec == nil {
		return nil, errGroupSecretsNotFound
	}
	raw, err := cs.decryptWithLabel(initKeyPriv, []byte("Welcome"), w.encryptedGroupInfo, sec.encryptedGroupSecrets.kemOutput, sec.encryptedGroupSecrets.ciphertext)
	if err != nil {
		return nil, err
	}
	var gs groupSecrets
	if err := unmarshalRaw(raw, &gs); err != nil {
		return nil, err
	}
	return &gs, nil
}

func (w *Welcome) decryptGroupInfo(joinerSecret, pskSecret []byte) (*groupInfo, error) {
	cs := w.cipherSuite
	welcomeSecret, err := extractWelcomeSecret(cs, joinerSecret, pskSecret)
	if err != nil {
		return nil, err
	}

	welcomeNonce, err := cs.expandWithLabel(welcomeSecret, []byte("nonce"), nil, uint16(cs.AEADNonceSize()))
	if err != nil {
		return nil, err
	}
	welcomeKey, err := cs.expandWithLabel(welcomeSecret, []byte("key"), nil, uint16(cs.AEADKeySize()))
	if err != nil {
		return nil, err
	}

	rawGroupInfo, err := cs.aeadOpen(welcomeKey, welcomeNonce, w.encryptedGroupInfo, nil)
	if err != nil {
		return nil, err
	}

	var info groupInfo
	if err := unmarshalRaw(rawGroupInfo, &info); err != nil {
		return nil, err
	}
	return &info, nil
}