package libsignal

/*
#include <signal_protocol.h>
#include <session_cipher.h>
#include <protocol.h>
*/
import "C"

import (
	"runtime"
	"unsafe"
)

// Ciphertext message types (protocol.h), discriminating what
// CiphertextMessage.Serialized holds.
const (
	// CiphertextSignalType is an ordinary ratchet message; a session must
	// already exist to decrypt one.
	CiphertextSignalType = C.CIPHERTEXT_SIGNAL_TYPE
	// CiphertextPreKeyType is a prekey message that also establishes the
	// session on the receiving end, via X3DH.
	CiphertextPreKeyType = C.CIPHERTEXT_PREKEY_TYPE
)

// CiphertextMessage is an encrypted message ready to be embedded in a
// stanza. Type distinguishes an ordinary ratchet message from a prekey
// message - the caller needs this to choose the right XEP-0384 wire
// representation (a plain <key> vs a <key prekey="true">).
type CiphertextMessage struct {
	Type       int
	Serialized []byte
}

// SessionCipher performs encrypt/decrypt for an already-established session.
type SessionCipher struct {
	raw  *C.session_cipher
	addr *cAddress
}

// NewSessionCipher creates a session cipher for storeCtx's local identity
// talking to remote. version selects the wire format for messages this
// cipher *produces* (ProtocolVersionLegacy or ProtocolVersionModern);
// decrypting adapts to whatever version the incoming message declares.
func NewSessionCipher(ctx *Context, storeCtx *StoreContext, remote Address, version int) (*SessionCipher, error) {
	addr := newCAddress(remote)

	var raw *C.session_cipher
	if code := C.session_cipher_create(&raw, storeCtx.raw, addr.ptr, ctx.raw); code != C.SG_SUCCESS {
		addr.free()
		return nil, newError("session_cipher_create", int(code))
	}
	if err := setCipherVersion(raw, version); err != nil {
		C.session_cipher_free(raw)
		addr.free()
		return nil, err
	}
	return &SessionCipher{raw: raw, addr: addr}, nil
}

// Close frees the session cipher. Do not use it afterward.
func (c *SessionCipher) Close() {
	if c.raw != nil {
		C.session_cipher_free(c.raw)
		c.raw = nil
	}
	c.addr.free()
}

// Encrypt encrypts plaintext for this cipher's session, in whichever
// protocol version the cipher was created with.
func (c *SessionCipher) Encrypt(plaintext []byte) (*CiphertextMessage, error) {
	var ptr *C.uint8_t
	if len(plaintext) > 0 {
		ptr = (*C.uint8_t)(unsafe.Pointer(&plaintext[0]))
	}
	var msg *C.ciphertext_message
	code := C.session_cipher_encrypt(c.raw, ptr, C.size_t(len(plaintext)), &msg)
	runtime.KeepAlive(plaintext)
	if code != C.SG_SUCCESS {
		return nil, newError("session_cipher_encrypt", int(code))
	}
	defer C.signal_type_unref((*C.signal_type_base)(unsafe.Pointer(msg)))

	return &CiphertextMessage{
		Type:       int(C.ciphertext_message_get_type(msg)),
		Serialized: bufferToBytes(C.ciphertext_message_get_serialized(msg)),
	}, nil
}

// Decrypt decrypts a message received from this cipher's remote peer.
//
// useOmemoFraming selects between the legacy wire encoding and the OMEMO
// (protocol v4) one for parsing msg.Serialized - it must match whichever
// namespace/version the sender actually used, not necessarily this
// cipher's own outgoing version (a peer's incoming and outgoing wire
// format need not match while a conversation is transitioning between
// legacy and modern OMEMO). ownRegistrationID is this identity's own
// registration id (IdentityStore.GetLocalRegistrationID) - only consulted
// for CiphertextPreKeyType messages under OMEMO framing, which the C API
// requires it for.
func (c *SessionCipher) Decrypt(ctx *Context, msg CiphertextMessage, useOmemoFraming bool, ownRegistrationID uint32) ([]byte, error) {
	var ptr *C.uint8_t
	if len(msg.Serialized) > 0 {
		ptr = (*C.uint8_t)(unsafe.Pointer(&msg.Serialized[0]))
	}

	var plaintextBuf *C.signal_buffer
	var code C.int

	switch msg.Type {
	case CiphertextPreKeyType:
		var preKeyMsg *C.pre_key_signal_message
		if useOmemoFraming {
			preKeyMsg, code = deserializePreKeySignalMessageOmemo(ctx, ptr, C.size_t(len(msg.Serialized)), ownRegistrationID)
		} else {
			code = C.pre_key_signal_message_deserialize(&preKeyMsg, ptr, C.size_t(len(msg.Serialized)), ctx.raw)
		}
		if code != C.SG_SUCCESS {
			break
		}
		defer C.signal_type_unref((*C.signal_type_base)(unsafe.Pointer(preKeyMsg)))
		code = C.session_cipher_decrypt_pre_key_signal_message(c.raw, preKeyMsg, nil, &plaintextBuf)
	case CiphertextSignalType:
		var signalMsg *C.signal_message
		if useOmemoFraming {
			signalMsg, code = deserializeSignalMessageOmemo(ctx, ptr, C.size_t(len(msg.Serialized)))
		} else {
			code = C.signal_message_deserialize(&signalMsg, ptr, C.size_t(len(msg.Serialized)), ctx.raw)
		}
		if code != C.SG_SUCCESS {
			break
		}
		defer C.signal_type_unref((*C.signal_type_base)(unsafe.Pointer(signalMsg)))
		code = C.session_cipher_decrypt_signal_message(c.raw, signalMsg, nil, &plaintextBuf)
	default:
		code = C.int(ErrInvalidMessage)
	}
	runtime.KeepAlive(msg.Serialized)

	if code != C.SG_SUCCESS {
		return nil, newError("session_cipher_decrypt", int(code))
	}
	defer freeBuffer(plaintextBuf)
	return bufferToBytes(plaintextBuf), nil
}