// 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"

	"dev.narayana.im/narayana/telegabber/e2ee/omemo/libsignal"
)

// 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
// omemo.go's job, built on top of EncryptOuter/DecryptOuter here.
//
// Two schemes, matching the two wire versions (confirmed against the
// XEP-0384 spec text directly, not memory - see the M2 kickoff discussion):
//
//   - Legacy (ProtocolVersionLegacy, 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.
//
//   - Modern (ProtocolVersionModern, OMEMO 1/2): the plaintext is first
//     wrapped in a minimal XEP-0420 SCE envelope (sce.go). 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. 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 (
	legacyKeyLen = 16 // AES-128
	legacyIVLen  = 12 // GCM nonce, standard size (crypto/cipher.NewGCM's default)
	legacyTagLen = 16

	modernKeyLen     = 32 // AES-256
	modernAuthKeyLen = 32
	modernIVLen      = 16 // AES block size
	modernTagLen     = 16
	modernHKDFInfo   = "OMEMO Payload"
)

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

// EncryptOuter performs OMEMO's outer content encryption for plaintext,
// per version (ProtocolVersionLegacy or ProtocolVersionModern). fromJID is
// used as the SCE <from> binding for modern; ignored for legacy.
func EncryptOuter(version int, plaintext []byte, fromJID string) (*EncryptedPayload, error) {
	switch version {
	case libsignal.ProtocolVersionLegacy:
		ciphertext, iv, contentKey, err := encryptPayloadLegacy(plaintext)
		if err != nil {
			return nil, err
		}
		return &EncryptedPayload{Ciphertext: ciphertext, IV: iv, ContentKey: contentKey}, nil
	case libsignal.ProtocolVersionModern:
		sceXML, err := sceEncode(plaintext, fromJID)
		if err != nil {
			return nil, err
		}
		ciphertext, contentKey, err := encryptPayloadModern(sceXML)
		if err != nil {
			return nil, err
		}
		return &EncryptedPayload{Ciphertext: ciphertext, ContentKey: contentKey}, nil
	default:
		return nil, fmt.Errorf("omemo: unsupported protocol version %d", version)
	}
}

// DecryptOuter reverses EncryptOuter.
func DecryptOuter(version int, payload *EncryptedPayload) ([]byte, error) {
	switch version {
	case libsignal.ProtocolVersionLegacy:
		return decryptPayloadLegacy(payload.Ciphertext, payload.IV, payload.ContentKey)
	case libsignal.ProtocolVersionModern:
		sceXML, err := decryptPayloadModern(payload.Ciphertext, payload.ContentKey)
		if err != nil {
			return nil, err
		}
		return sceDecode(sceXML)
	default:
		return nil, fmt.Errorf("omemo: unsupported protocol version %d", version)
	}
}

func encryptPayloadLegacy(plaintext []byte) (ciphertext, iv, contentKey []byte, err error) {
	key := make([]byte, legacyKeyLen)
	if _, err := rand.Read(key); err != nil {
		return nil, nil, nil, err
	}
	iv = make([]byte, legacyIVLen)
	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) < legacyTagLen {
		return nil, nil, nil, errors.New("omemo: GCM output shorter than its own tag")
	}
	ciphertext = sealed[:len(sealed)-legacyTagLen]
	tag := sealed[len(sealed)-legacyTagLen:]
	contentKey = append(append([]byte{}, key...), tag...)
	return ciphertext, iv, contentKey, nil
}

func decryptPayloadLegacy(ciphertext, iv, contentKey []byte) ([]byte, error) {
	if len(contentKey) != legacyKeyLen+legacyTagLen {
		return nil, fmt.Errorf("omemo: legacy content key must be %d bytes, got %d", legacyKeyLen+legacyTagLen, len(contentKey))
	}
	key := contentKey[:legacyKeyLen]
	tag := contentKey[legacyKeyLen:]
	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(modernHKDFInfo))
	out := make([]byte, modernKeyLen+modernAuthKeyLen+modernIVLen)
	if _, err := io.ReadFull(h, out); err != nil {
		return nil, nil, nil, err
	}
	return out[:modernKeyLen], out[modernKeyLen : modernKeyLen+modernAuthKeyLen], out[modernKeyLen+modernAuthKeyLen:], nil
}

func encryptPayloadModern(plaintext []byte) (ciphertext, contentKey []byte, err error) {
	key := make([]byte, modernKeyLen)
	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)[:modernTagLen]

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

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

	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)[:modernTagLen]
	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=".."/> (legacy) or <key rid=".." kex=".."/>
// (modern) 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 (see the M2 architecture note: libomemo-c itself produces 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 (legacy vs OMEMO 2/SCE) 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   int
	SenderSID uint32
	IV        []byte // legacy only
	Payload   []byte // outer ciphertext
	Keys      []WireKey
}

// Encode serializes a WireEnvelope for e2ee.Envelope.Raw. This is
// Go-to-Go-only plumbing (produced by omemo.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
}