package mls

// MLS framing crypto operations (RFC 9420 §6).
// Sign/verify/encrypt/decrypt methods for message framing.

import "smesh.lol/web/common/jsbridge/subtle"

// --- Signing ---

func signAuthenticatedContent(cs CipherSuite, signKey signaturePrivateKey, wf wireFormat, content *framedContent, ctx *groupContext) (*authenticatedContent, error) {
	authContent := authenticatedContent{
		wireFormat: wf,
		content:    *content,
	}
	tbs := authContent.framedContentTBS(ctx)
	sig, err := signFramedContent(cs, signKey, tbs)
	if err != nil {
		return nil, err
	}
	authContent.auth.signature = sig
	return &authContent, nil
}

func signFramedContent(cs CipherSuite, signKey signaturePrivateKey, content *framedContentTBS) ([]byte, error) {
	raw, err := marshalRaw(content)
	if err != nil {
		return nil, err
	}
	return cs.signWithLabel(signKey, []byte("FramedContentTBS"), raw)
}

func (authContent *authenticatedContent) verifySignature(verifKey signaturePublicKey, ctx *groupContext) bool {
	raw, err := marshalRaw(authContent.framedContentTBS(ctx))
	if err != nil {
		return false
	}
	return ctx.cipherSuite.verifyWithLabel(verifKey, []byte("FramedContentTBS"), raw, authContent.auth.signature)
}

func (authContent *authenticatedContent) generateProposalRef(cs CipherSuite) (proposalRef, error) {
	if authContent.content.contentType != contentTypeProposal {
		panic("mls: not a proposal")
	}
	raw, err := marshalRaw(authContent)
	if err != nil {
		return nil, err
	}
	h, err := cs.refHash([]byte("MLS 1.0 Proposal Reference"), raw)
	if err != nil {
		return nil, err
	}
	return proposalRef(h), nil
}

// --- FramedContentAuthData ---

func (authData *framedContentAuthData) verifyConfirmationTag(cs CipherSuite, confirmationKey, confirmedTranscriptHash []byte) bool {
	if len(authData.confirmationTag) == 0 {
		return false
	}
	return cs.verifyMAC(confirmationKey, confirmedTranscriptHash, authData.confirmationTag)
}

func (authData *framedContentAuthData) verifySignature(cs CipherSuite, verifKey signaturePublicKey, content *framedContentTBS) bool {
	raw, err := marshalRaw(content)
	if err != nil {
		return false
	}
	return cs.verifyWithLabel(verifKey, []byte("FramedContentTBS"), raw, authData.signature)
}

// --- PublicMessage ---

func signPublicMessage(cs CipherSuite, signKey signaturePrivateKey, content *framedContent, ctx *groupContext) (*publicMessage, error) {
	authContent, err := signAuthenticatedContent(cs, signKey, wireFormatMLSPublicMessage, content, ctx)
	if err != nil {
		return nil, err
	}
	return &publicMessage{
		content: authContent.content,
		auth:    authContent.auth,
	}, nil
}

func (msg *publicMessage) signMembershipTag(cs CipherSuite, membershipKey []byte, ctx *groupContext) error {
	if msg.content.sender.senderType != senderTypeMember {
		return nil
	}
	raw, err := marshalRaw(msg.authenticatedContentTBM(ctx))
	if err != nil {
		return err
	}
	msg.membershipTag = cs.signMAC(membershipKey, raw)
	return nil
}

func (msg *publicMessage) verifyMembershipTag(membershipKey []byte, ctx *groupContext) bool {
	if msg.content.sender.senderType != senderTypeMember {
		return true
	}
	raw, err := marshalRaw(msg.authenticatedContentTBM(ctx))
	if err != nil {
		return false
	}
	return ctx.cipherSuite.verifyMAC(membershipKey, raw, msg.membershipTag)
}

// --- PrivateMessage ---

func newSenderData(li leafIndex, generation uint32) (*senderData, error) {
	data := senderData{
		leafIndex:  li,
		generation: generation,
	}
	var rg [4]byte
	subtle.RandomBytes(rg[:])
	data.reuseGuard = rg
	return &data, nil
}

func sampleCiphertext(cs CipherSuite, ciphertext []byte) []byte {
	n := cs.ExtractSize()
	if len(ciphertext) < n {
		return ciphertext
	}
	return ciphertext[:n]
}

func expandSenderDataKey(cs CipherSuite, senderDataSecret, ciphertext []byte) ([]byte, error) {
	sample := sampleCiphertext(cs, ciphertext)
	return cs.expandWithLabel(senderDataSecret, []byte("key"), sample, uint16(cs.AEADKeySize()))
}

func expandSenderDataNonce(cs CipherSuite, senderDataSecret, ciphertext []byte) ([]byte, error) {
	sample := sampleCiphertext(cs, ciphertext)
	return cs.expandWithLabel(senderDataSecret, []byte("nonce"), sample, uint16(cs.AEADNonceSize()))
}

func derivePrivateMessageKeyAndNonce(cs CipherSuite, secret ratchetSecret, reuseGuard [4]byte) (key, nonce []byte, err error) {
	key, err = secret.deriveKey(cs)
	if err != nil {
		return nil, nil, err
	}
	nonce, err = secret.deriveNonce(cs)
	if err != nil {
		return nil, nil, err
	}
	for i := range reuseGuard {
		nonce[i] = nonce[i] ^ reuseGuard[i]
	}
	return key, nonce, nil
}

func signPrivateMessageContent(cs CipherSuite, signKey signaturePrivateKey, content *framedContent, ctx *groupContext) (*privateMessageContent, error) {
	authContent, err := signAuthenticatedContent(cs, signKey, wireFormatMLSPrivateMessage, content, ctx)
	if err != nil {
		return nil, err
	}
	return &privateMessageContent{
		applicationData: content.applicationData,
		proposal:        content.proposal,
		commit:          content.commit,
		auth:            authContent.auth,
	}, nil
}

func encryptPrivateMessageContent(cs CipherSuite, secret ratchetSecret, content *framedContent, privContent *privateMessageContent, reuseGuard [4]byte) ([]byte, error) {
	var w Writer
	privContent.marshal(&w, content.contentType)
	plaintext, err := w.bytes()
	if err != nil {
		return nil, err
	}

	key, nonce, err := derivePrivateMessageKeyAndNonce(cs, secret, reuseGuard)
	if err != nil {
		return nil, err
	}

	aad := privateContentAAD{
		groupID:           content.groupID,
		epoch:             content.epoch,
		contentType:       content.contentType,
		authenticatedData: content.authenticatedData,
	}
	rawAAD, err := marshalRaw(&aad)
	if err != nil {
		return nil, err
	}

	return cs.aeadSeal(key, nonce, plaintext, rawAAD)
}

func encryptSenderData(cs CipherSuite, senderDataSecret []byte, sd *senderData, content *framedContent, ciphertext []byte) ([]byte, error) {
	key, err := expandSenderDataKey(cs, senderDataSecret, ciphertext)
	if err != nil {
		return nil, err
	}
	nonce, err := expandSenderDataNonce(cs, senderDataSecret, ciphertext)
	if err != nil {
		return nil, err
	}

	aad := senderDataAAD{
		groupID:     content.groupID,
		epoch:       content.epoch,
		contentType: content.contentType,
	}
	rawAAD, err := marshalRaw(&aad)
	if err != nil {
		return nil, err
	}

	rawSD, err := marshalRaw(sd)
	if err != nil {
		return nil, err
	}

	return cs.aeadSeal(key, nonce, rawSD, rawAAD)
}

func encryptPrivateMessage(cs CipherSuite, secret ratchetSecret, senderDataSecret []byte, content *framedContent, privContent *privateMessageContent, sd *senderData) (*privateMessage, error) {
	ct, err := encryptPrivateMessageContent(cs, secret, content, privContent, sd.reuseGuard)
	if err != nil {
		return nil, err
	}
	encSD, err := encryptSenderData(cs, senderDataSecret, sd, content, ct)
	if err != nil {
		return nil, err
	}
	return &privateMessage{
		groupID:             content.groupID,
		epoch:               content.epoch,
		contentType:         content.contentType,
		authenticatedData:   content.authenticatedData,
		encryptedSenderData: encSD,
		ciphertext:          ct,
	}, nil
}

func (msg *privateMessage) decryptSenderData(cs CipherSuite, senderDataSecret []byte) (*senderData, error) {
	key, err := expandSenderDataKey(cs, senderDataSecret, msg.ciphertext)
	if err != nil {
		return nil, err
	}
	nonce, err := expandSenderDataNonce(cs, senderDataSecret, msg.ciphertext)
	if err != nil {
		return nil, err
	}

	aad := senderDataAAD{
		groupID:     msg.groupID,
		epoch:       msg.epoch,
		contentType: msg.contentType,
	}
	rawAAD, err := marshalRaw(&aad)
	if err != nil {
		return nil, err
	}

	rawSD, err := cs.aeadOpen(key, nonce, msg.encryptedSenderData, rawAAD)
	if err != nil {
		return nil, err
	}

	var sd senderData
	if err := unmarshalRaw(rawSD, &sd); err != nil {
		return nil, err
	}
	return &sd, nil
}

func (msg *privateMessage) decryptContent(cs CipherSuite, secret ratchetSecret, reuseGuard [4]byte) (*privateMessageContent, error) {
	key, nonce, err := derivePrivateMessageKeyAndNonce(cs, secret, reuseGuard)
	if err != nil {
		return nil, err
	}

	aad := privateContentAAD{
		groupID:           msg.groupID,
		epoch:             msg.epoch,
		contentType:       msg.contentType,
		authenticatedData: msg.authenticatedData,
	}
	rawAAD, err := marshalRaw(&aad)
	if err != nil {
		return nil, err
	}

	rawContent, err := cs.aeadOpen(key, nonce, msg.ciphertext, rawAAD)
	if err != nil {
		return nil, err
	}

	r := newReader(rawContent)
	var content privateMessageContent
	if err := content.unmarshal(&r, msg.contentType); err != nil {
		return nil, err
	}

	// Verify remaining bytes are zero-padding
	for !r.empty() {
		b, _ := r.readByte()
		if b != 0 {
			return nil, errNonZeroPadding
		}
	}

	return &content, nil
}