package libsignal

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

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	rand2 "crypto/rand"
	"crypto/sha256"
	"crypto/sha512"
	"errors"
	"hash"
	"runtime/cgo"
	"unsafe"
)

// libomemo-c ships no default signal_crypto_provider (confirmed by reading
// signal_protocol.h - only test-only OpenSSL/CommonCrypto implementations
// exist, never installed). This file implements the five callback groups
// documented there - random, HMAC-SHA256, SHA-512, and AES encrypt/decrypt
// in the two modes the library actually asks for (CTR-nopadding and
// CBC-PKCS5; no GCM) - entirely against Go's standard library.
//
// hmac_context/digest_context are opaque per-operation state the library
// holds between init/update/final/cleanup calls. Since C can't safely hold
// a long-lived Go pointer, each is a runtime/cgo.Handle round-tripped
// through a uintptr, exactly like the store callbacks in callbacks.go.

//export go_crypto_random
func go_crypto_random(data *C.uint8_t, length C.size_t, userData unsafe.Pointer) C.int {
	buf := unsafe.Slice((*byte)(unsafe.Pointer(data)), int(length))
	if _, err := rand2.Read(buf); err != nil {
		return C.int(ErrUnknown)
	}
	return C.SG_SUCCESS
}

type hmacState struct {
	h hash.Hash
}

//export go_crypto_hmac_sha256_init
func go_crypto_hmac_sha256_init(hmacContext *unsafe.Pointer, key *C.uint8_t, keyLen C.size_t, userData unsafe.Pointer) C.int {
	keyBytes := cBytesToGo(key, keyLen)
	handle := cgo.NewHandle(&hmacState{h: hmac.New(sha256.New, keyBytes)})
	// go vet flags this as "possible misuse of unsafe.Pointer" - expected: a
	// cgo.Handle is an opaque uintptr token, never dereferenced as a real
	// address, and signal_protocol.h's callback shape requires void* here.
	*hmacContext = unsafe.Pointer(uintptr(handle))
	return C.SG_SUCCESS
}

//export go_crypto_hmac_sha256_update
func go_crypto_hmac_sha256_update(hmacContext unsafe.Pointer, data *C.uint8_t, dataLen C.size_t, userData unsafe.Pointer) C.int {
	st, ok := cgo.Handle(uintptr(hmacContext)).Value().(*hmacState)
	if !ok {
		return C.int(ErrUnknown)
	}
	st.h.Write(cBytesToGo(data, dataLen))
	return C.SG_SUCCESS
}

//export go_crypto_hmac_sha256_final
func go_crypto_hmac_sha256_final(hmacContext unsafe.Pointer, output **C.signal_buffer, userData unsafe.Pointer) C.int {
	st, ok := cgo.Handle(uintptr(hmacContext)).Value().(*hmacState)
	if !ok {
		return C.int(ErrUnknown)
	}
	*output = bytesToBuffer(st.h.Sum(nil))
	st.h.Reset()
	return C.SG_SUCCESS
}

//export go_crypto_hmac_sha256_cleanup
func go_crypto_hmac_sha256_cleanup(hmacContext unsafe.Pointer, userData unsafe.Pointer) {
	cgo.Handle(uintptr(hmacContext)).Delete()
}

type digestState struct {
	h hash.Hash
}

//export go_crypto_sha512_init
func go_crypto_sha512_init(digestContext *unsafe.Pointer, userData unsafe.Pointer) C.int {
	handle := cgo.NewHandle(&digestState{h: sha512.New()})
	// See the identical vet note in go_crypto_hmac_sha256_init above.
	*digestContext = unsafe.Pointer(uintptr(handle))
	return C.SG_SUCCESS
}

//export go_crypto_sha512_update
func go_crypto_sha512_update(digestContext unsafe.Pointer, data *C.uint8_t, dataLen C.size_t, userData unsafe.Pointer) C.int {
	st, ok := cgo.Handle(uintptr(digestContext)).Value().(*digestState)
	if !ok {
		return C.int(ErrUnknown)
	}
	st.h.Write(cBytesToGo(data, dataLen))
	return C.SG_SUCCESS
}

//export go_crypto_sha512_final
func go_crypto_sha512_final(digestContext unsafe.Pointer, output **C.signal_buffer, userData unsafe.Pointer) C.int {
	st, ok := cgo.Handle(uintptr(digestContext)).Value().(*digestState)
	if !ok {
		return C.int(ErrUnknown)
	}
	*output = bytesToBuffer(st.h.Sum(nil))
	st.h.Reset()
	return C.SG_SUCCESS
}

//export go_crypto_sha512_cleanup
func go_crypto_sha512_cleanup(digestContext unsafe.Pointer, userData unsafe.Pointer) {
	cgo.Handle(uintptr(digestContext)).Delete()
}

//export go_crypto_encrypt
func go_crypto_encrypt(output **C.signal_buffer, cipherMode C.int, key *C.uint8_t, keyLen C.size_t, iv *C.uint8_t, ivLen C.size_t, plaintext *C.uint8_t, plaintextLen C.size_t, userData unsafe.Pointer) C.int {
	block, err := aes.NewCipher(cBytesToGo(key, keyLen))
	if err != nil {
		return C.int(ErrUnknown)
	}
	ivBytes := cBytesToGo(iv, ivLen)
	plaintextBytes := cBytesToGo(plaintext, plaintextLen)

	var out []byte
	switch cipherMode {
	case C.SG_CIPHER_AES_CTR_NOPADDING:
		out = make([]byte, len(plaintextBytes))
		cipher.NewCTR(block, ivBytes).XORKeyStream(out, plaintextBytes)
	case C.SG_CIPHER_AES_CBC_PKCS5:
		padded := pkcs7Pad(plaintextBytes, aes.BlockSize)
		out = make([]byte, len(padded))
		cipher.NewCBCEncrypter(block, ivBytes).CryptBlocks(out, padded)
	default:
		return C.int(ErrInvalidArgument)
	}

	*output = bytesToBuffer(out)
	return C.SG_SUCCESS
}

//export go_crypto_decrypt
func go_crypto_decrypt(output **C.signal_buffer, cipherMode C.int, key *C.uint8_t, keyLen C.size_t, iv *C.uint8_t, ivLen C.size_t, ciphertext *C.uint8_t, ciphertextLen C.size_t, userData unsafe.Pointer) C.int {
	block, err := aes.NewCipher(cBytesToGo(key, keyLen))
	if err != nil {
		return C.int(ErrUnknown)
	}
	ivBytes := cBytesToGo(iv, ivLen)
	ciphertextBytes := cBytesToGo(ciphertext, ciphertextLen)

	var out []byte
	switch cipherMode {
	case C.SG_CIPHER_AES_CTR_NOPADDING:
		out = make([]byte, len(ciphertextBytes))
		cipher.NewCTR(block, ivBytes).XORKeyStream(out, ciphertextBytes)
	case C.SG_CIPHER_AES_CBC_PKCS5:
		if len(ciphertextBytes) == 0 || len(ciphertextBytes)%aes.BlockSize != 0 {
			return C.int(ErrInvalidMessage)
		}
		decrypted := make([]byte, len(ciphertextBytes))
		cipher.NewCBCDecrypter(block, ivBytes).CryptBlocks(decrypted, ciphertextBytes)
		unpadded, err := pkcs7Unpad(decrypted)
		if err != nil {
			return C.int(ErrInvalidMessage)
		}
		out = unpadded
	default:
		return C.int(ErrInvalidArgument)
	}

	*output = bytesToBuffer(out)
	return C.SG_SUCCESS
}

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("libsignal: cannot unpad empty data")
	}
	padLen := int(data[n-1])
	if padLen == 0 || padLen > n || padLen > aes.BlockSize {
		return nil, errors.New("libsignal: invalid PKCS7 padding")
	}
	for _, b := range data[n-padLen:] {
		if int(b) != padLen {
			return nil, errors.New("libsignal: invalid PKCS7 padding")
		}
	}
	return data[:n-padLen], nil
}