// Package omemo is the OMEMO (XEP-0384) implementation of e2ee.Backend.
package omemo

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/gob"
	"errors"
	"fmt"
	"io"

	"golang.org/x/crypto/hkdf"
)

// This file is OMEMO's outer, non-Signal content encryption: the layer
// that encrypts the actual message body once per message, independent of
// how many recipient devices it goes to. The resulting small "content key"
// blob (not the body itself) is what gets Signal/Double-Ratchet-encrypted
// once per recipient device (see libsignal.SessionCipher.Encrypt) - that
// per-device wrapping, and the device/session bookkeeping around it, is
// identity.go/encrypt.go's job, built on top of EncryptOuter/DecryptOuter
// here.
//
// Three versions (see version.go), confirmed against the XEP-0384 spec
// text and its version history directly, not memory:
//
//   - Omemo0 (eu.siacs.conversations.axolotl): AES-128-GCM directly on the
//     plaintext body. The GCM tag is stripped from the ciphertext and
//     instead appended to the key (key||tag, 32 bytes) before that blob is
//     Signal-encrypted per recipient; <payload> carries ciphertext only,
//     <iv> carries the GCM nonce. This is what we (and current real
//     clients) produce; decryptPayloadOmemo0 additionally accepts an
//     older, non-conformant 16-byte-key-only layout some real-world
//     senders still use (see its own doc comment) - real-world interop,
//     confirmed against actual client behavior/history, not assumed.
//
//   - Omemo1 and Omemo2 share the same payload crypto (introduced at
//     Omemo1/XEP-0384 0.4.0, unchanged since): the plaintext is first
//     wrapped in an XEP-0420 SCE envelope - Omemo1 and Omemo2 use two
//     different SCE dialects, see sce.go - then that envelope is encrypted
//     with AES-256-CBC using a key/authKey/iv derived via
//     HKDF-SHA256(random 32-byte key, salt=32 zero bytes, info="OMEMO
//     Payload") -> 80 bytes, split 32/32/16. The ciphertext is
//     HMAC-SHA256'd with authKey and truncated to 16 bytes (confirmed
//     against the current spec text directly - an earlier spec version
//     had a 32-vs-16-byte inconsistency here, fixed in 0.8.2). key||tag
//     (48 bytes) is the per-recipient blob; <payload> carries ciphertext
//     only; there is no <iv> element - the receiver re-derives it from the
//     key via the same HKDF.

const (
	omemo0KeyLen = 16 // AES-128
	omemo0IVLen  = 12 // GCM nonce, standard size (crypto/cipher.NewGCM's default)
	omemo0TagLen = 16

	sceKeyLen     = 32 // AES-256 (Omemo1/Omemo2 shared payload crypto)
	sceAuthKeyLen = 32
	sceIVLen      = 16 // AES block size
	sceTagLen     = 16
	sceHKDFInfo   = "OMEMO Payload"
)

// EncryptedPayload is the result of OMEMO's outer content-encryption step.
type EncryptedPayload struct {
	Ciphertext []byte
	IV         []byte // Omemo0 only; nil for Omemo1/Omemo2
	ContentKey []byte // the blob to Signal-encrypt once per recipient device
}

// EncryptOuter performs OMEMO's outer content encryption for plaintext,
// per version. fromJID is used as the SCE <from> binding for Omemo1/
// Omemo2; ignored for Omemo0 (which has no SCE wrapping at all).
func EncryptOuter(version Version, plaintext []byte, fromJID string) (*EncryptedPayload, error) {
	switch version {
	case Omemo0:
		ciphertext, iv, contentKey, err := encryptPayloadOmemo0(plaintext)
		if err != nil {
			return nil, err
		}
		return &EncryptedPayload{Ciphertext: ciphertext, IV: iv, ContentKey: contentKey}, nil
	case Omemo1:
		sceXML, err := sceEncodeV0(plaintext, fromJID)
		if err != nil {
			return nil, err
		}
		ciphertext, contentKey, err := encryptPayloadSCE(sceXML)
		if err != nil {
			return nil, err
		}
		return &EncryptedPayload{Ciphertext: ciphertext, ContentKey: contentKey}, nil
	case Omemo2:
		sceXML, err := sceEncodeV1(plaintext, fromJID)
		if err != nil {
			return nil, err
		}
		ciphertext, contentKey, err := encryptPayloadSCE(sceXML)
		if err != nil {
			return nil, err
		}
		return &EncryptedPayload{Ciphertext: ciphertext, ContentKey: contentKey}, nil
	default:
		return nil, fmt.Errorf("omemo: EncryptOuter: unsupported version %v", version)
	}
}

// DecryptOuter reverses EncryptOuter.
func DecryptOuter(version Version, payload *EncryptedPayload) ([]byte, error) {
	switch version {
	case Omemo0:
		return decryptPayloadOmemo0(payload.Ciphertext, payload.IV, payload.ContentKey)
	case Omemo1:
		sceXML, err := decryptPayloadSCE(payload.Ciphertext, payload.ContentKey)
		if err != nil {
			return nil, err
		}
		return sceDecodeV0(sceXML)
	case Omemo2:
		sceXML, err := decryptPayloadSCE(payload.Ciphertext, payload.ContentKey)
		if err != nil {
			return nil, err
		}
		return sceDecodeV1(sceXML)
	default:
		return nil, fmt.Errorf("omemo: DecryptOuter: unsupported version %v", version)
	}
}

func encryptPayloadOmemo0(plaintext []byte) (ciphertext, iv, contentKey []byte, err error) {
	key := make([]byte, omemo0KeyLen)
	if _, err := rand.Read(key); err != nil {
		return nil, nil, nil, err
	}
	iv = make([]byte, omemo0IVLen)
	if _, err := rand.Read(iv); err != nil {
		return nil, nil, nil, err
	}
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, nil, nil, err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, nil, nil, err
	}
	sealed := gcm.Seal(nil, iv, plaintext, nil)
	if len(sealed) < omemo0TagLen {
		return nil, nil, nil, errors.New("omemo: GCM output shorter than its own tag")
	}
	ciphertext = sealed[:len(sealed)-omemo0TagLen]
	tag := sealed[len(sealed)-omemo0TagLen:]
	contentKey = append(append([]byte{}, key...), tag...)
	return ciphertext, iv, contentKey, nil
}

// decryptPayloadOmemo0 accepts two real-world wire layouts for the
// per-device content key, dispatching on its length:
//
//   - 32 bytes (omemo0KeyLen+omemo0TagLen): the spec-correct, current
//     layout this package also produces (see encryptPayloadOmemo0) - key
//     and GCM tag concatenated in <key>, <payload> is ciphertext only.
//     Confirmed as what current real clients use (python-oldmemo, the
//     backend behind Gajim's OMEMO plugin and slixmpp-omemo; Profanity;
//     the xmpp-parsers crate's doc comment) - not assumed.
//   - 16 bytes (omemo0KeyLen): an older, non-conformant layout some
//     real-world senders still use, where <key> is just the raw AES key
//     and the GCM tag instead rides along appended to <payload>
//     (ciphertext||tag - the natural output of a plain AEAD Seal call).
//     Confirmed via Conversations' own git history (commit e38a9cd729bf,
//     "no longer accept auth tag appended to payload", ~2020-01 - meaning
//     Conversations itself accepted exactly this for years before that)
//     and Dino's dino#303 (a bug from mishandling this same split) - not
//     a hypothetical. Accepted here defensively for interop, the same way
//     Conversations itself did.
func decryptPayloadOmemo0(ciphertext, iv, contentKey []byte) ([]byte, error) {
	switch len(contentKey) {
	case omemo0KeyLen + omemo0TagLen:
		return openOmemo0GCM(ciphertext, iv, contentKey[:omemo0KeyLen], contentKey[omemo0KeyLen:])
	case omemo0KeyLen:
		if len(ciphertext) < omemo0TagLen {
			return nil, fmt.Errorf("omemo: omemo0 payload too short to contain a trailing GCM tag: %d bytes", len(ciphertext))
		}
		tag := ciphertext[len(ciphertext)-omemo0TagLen:]
		return openOmemo0GCM(ciphertext[:len(ciphertext)-omemo0TagLen], iv, contentKey, tag)
	default:
		return nil, fmt.Errorf("omemo: omemo0 content key must be %d or %d bytes, got %d", omemo0KeyLen, omemo0KeyLen+omemo0TagLen, len(contentKey))
	}
}

func openOmemo0GCM(ciphertext, iv, key, tag []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	sealed := append(append([]byte{}, ciphertext...), tag...)
	return gcm.Open(nil, iv, sealed, nil)
}

func hkdfSplit(key []byte) (encKey, authKey, iv []byte, err error) {
	h := hkdf.New(sha256.New, key, make([]byte, sha256.Size), []byte(sceHKDFInfo))
	out := make([]byte, sceKeyLen+sceAuthKeyLen+sceIVLen)
	if _, err := io.ReadFull(h, out); err != nil {
		return nil, nil, nil, err
	}
	return out[:sceKeyLen], out[sceKeyLen : sceKeyLen+sceAuthKeyLen], out[sceKeyLen+sceAuthKeyLen:], nil
}

// encryptPayloadSCE/decryptPayloadSCE are the payload crypto shared by
// Omemo1 and Omemo2 - only the SCE dialect wrapping plaintext before this
// runs differs between them (see sce.go).
func encryptPayloadSCE(plaintext []byte) (ciphertext, contentKey []byte, err error) {
	key := make([]byte, sceKeyLen)
	if _, err := rand.Read(key); err != nil {
		return nil, nil, err
	}
	encKey, authKey, iv, err := hkdfSplit(key)
	if err != nil {
		return nil, nil, err
	}
	block, err := aes.NewCipher(encKey)
	if err != nil {
		return nil, nil, err
	}
	padded := pkcs7Pad(plaintext, aes.BlockSize)
	ciphertext = make([]byte, len(padded))
	cipher.NewCBCEncrypter(block, iv).CryptBlocks(ciphertext, padded)

	mac := hmac.New(sha256.New, authKey)
	mac.Write(ciphertext)
	tag := mac.Sum(nil)[:sceTagLen]

	contentKey = append(append([]byte{}, key...), tag...)
	return ciphertext, contentKey, nil
}

func decryptPayloadSCE(ciphertext, contentKey []byte) ([]byte, error) {
	if len(contentKey) != sceKeyLen+sceTagLen {
		return nil, fmt.Errorf("omemo: SCE content key must be %d bytes, got %d", sceKeyLen+sceTagLen, len(contentKey))
	}
	key := contentKey[:sceKeyLen]
	tag := contentKey[sceKeyLen:]

	encKey, authKey, iv, err := hkdfSplit(key)
	if err != nil {
		return nil, err
	}

	mac := hmac.New(sha256.New, authKey)
	mac.Write(ciphertext)
	expectedTag := mac.Sum(nil)[:sceTagLen]
	if !hmac.Equal(expectedTag, tag) {
		return nil, errors.New("omemo: payload authentication failed")
	}

	if len(ciphertext) == 0 || len(ciphertext)%aes.BlockSize != 0 {
		return nil, errors.New("omemo: invalid ciphertext length")
	}
	block, err := aes.NewCipher(encKey)
	if err != nil {
		return nil, err
	}
	padded := make([]byte, len(ciphertext))
	cipher.NewCBCDecrypter(block, iv).CryptBlocks(padded, ciphertext)
	return pkcs7Unpad(padded)
}

func pkcs7Pad(data []byte, blockSize int) []byte {
	padLen := blockSize - (len(data) % blockSize)
	padded := make([]byte, len(data)+padLen)
	copy(padded, data)
	for i := len(data); i < len(padded); i++ {
		padded[i] = byte(padLen)
	}
	return padded
}

func pkcs7Unpad(data []byte) ([]byte, error) {
	n := len(data)
	if n == 0 {
		return nil, errors.New("omemo: cannot unpad empty data")
	}
	padLen := int(data[n-1])
	if padLen == 0 || padLen > n || padLen > aes.BlockSize {
		return nil, errors.New("omemo: invalid PKCS7 padding")
	}
	for _, b := range data[n-padLen:] {
		if int(b) != padLen {
			return nil, errors.New("omemo: invalid PKCS7 padding")
		}
	}
	return data[:n-padLen], nil
}

// WireKey is one recipient device's Signal-encrypted content-key blob,
// version-and-namespace-agnostic - the xmpp-layer glue turns this into a
// <key rid=".." prekey=".."/> (Omemo0) or <key rid=".." kex=".."/>
// (Omemo1/Omemo2) element; both attributes mean the same thing (IsPreKey).
type WireKey struct {
	RecipientJID string
	DeviceID     uint32
	IsPreKey     bool
	Ciphertext   []byte // libsignal.CiphertextMessage.Serialized - already in the correct per-version wire format (libomemo-c itself produces the OMEMOKeyExchange/OMEMOAuthenticatedMessage protobuf bytes for protocol v4, no extra wrapping needed here)
}

// WireEnvelope is the decoded structure of an OMEMO <encrypted> element,
// independent of which XML shape it came from or will be marshaled to -
// that XML-specific glue lives in the xmpp package, which imports this
// package explicitly when the omemo backend is active.
type WireEnvelope struct {
	Version   Version
	SenderSID uint32
	IV        []byte // Omemo0 only
	Payload   []byte // outer ciphertext
	Keys      []WireKey
}

// Encode serializes a WireEnvelope for e2ee.Envelope.Raw. This is
// Go-to-Go-only plumbing (produced by identity.go/encrypt.go's Encrypt,
// consumed by the xmpp-layer glue in the same process), not a wire format
// in its own right, so gob is a fine fit - no schema stability concerns.
func (e *WireEnvelope) Encode() ([]byte, error) {
	var buf bytes.Buffer
	if err := gob.NewEncoder(&buf).Encode(e); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

// DecodeWireEnvelope reverses WireEnvelope.Encode.
func DecodeWireEnvelope(raw []byte) (*WireEnvelope, error) {
	var e WireEnvelope
	if err := gob.NewDecoder(bytes.NewReader(raw)).Decode(&e); err != nil {
		return nil, err
	}
	return &e, nil
}