/*
 * Unit tests for validity pairing helper functions
 *
 * Copyright (C) 2024 libfprint contributors
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 */

#include <glib.h>
#include <string.h>

#include <openssl/evp.h>
#include <openssl/hmac.h>
#include <openssl/rand.h>
#include <openssl/core_names.h>
#include <openssl/param_build.h>
#include <openssl/ec.h>

#include "fpi-byte-utils.h"
#include "drivers/validity/validity.h"
#include "drivers/validity/validity_pair.h"
#include "drivers/validity/validity_hal.h"

/* ================================================================
 * T7.11: parse_flash_info — valid response
 * ================================================================ */
static void
test_parse_flash_info_valid (void)
{
  /* CMD 0x3e response format (after 2-byte status, already stripped):
   * [jid0:2LE][jid1:2LE][blocks:2LE][unknown0:2LE][blocksize:2LE]
   * [unknown1:2LE][pcnt:2LE] = 14 bytes minimum */
  guint8 data[14];

  memset (data, 0, sizeof (data));

  /* jid0=0x01, jid1=0x02, blocks=0x1000, unknown0=0, blocksize=0x100,
   * unknown1=0, pcnt=5 (5 partitions = device already paired) */
  FP_WRITE_UINT16_LE (data + 0, 0x0001); /* jid0 */
  FP_WRITE_UINT16_LE (data + 2, 0x0002); /* jid1 */
  FP_WRITE_UINT16_LE (data + 4, 0x1000); /* blocks */
  FP_WRITE_UINT16_LE (data + 6, 0x0000); /* unknown0 */
  FP_WRITE_UINT16_LE (data + 8, 0x0100); /* blocksize */
  FP_WRITE_UINT16_LE (data + 10, 0x0000); /* unknown1 */
  FP_WRITE_UINT16_LE (data + 12, 5);      /* pcnt = 5 */

  ValidityFlashIcParams ic;
  guint16 num_partitions;

  gboolean ok = validity_pair_parse_flash_info (data, sizeof (data),
                                                &ic, &num_partitions);
  g_assert_true (ok);
  g_assert_cmpuint (num_partitions, ==, 5);
  g_assert_cmpuint (ic.size, ==, 0x1000 * 0x0100);
  g_assert_cmpuint (ic.sector_size, ==, 0x1000);
  g_assert_cmpuint (ic.sector_erase_cmd, ==, 0x20);
}

/* ================================================================
 * T7.12: parse_flash_info — zero partitions means needs pairing
 * ================================================================ */
static void
test_parse_flash_info_needs_pairing (void)
{
  guint8 data[14];

  memset (data, 0, sizeof (data));

  FP_WRITE_UINT16_LE (data + 0, 0x0001);
  FP_WRITE_UINT16_LE (data + 2, 0x0002);
  FP_WRITE_UINT16_LE (data + 4, 0x0800);
  FP_WRITE_UINT16_LE (data + 8, 0x0200);
  FP_WRITE_UINT16_LE (data + 12, 0);  /* 0 partitions */

  ValidityFlashIcParams ic;
  guint16 num_partitions;

  gboolean ok = validity_pair_parse_flash_info (data, sizeof (data),
                                                &ic, &num_partitions);
  g_assert_true (ok);
  g_assert_cmpuint (num_partitions, ==, 0);
}

/* ================================================================
 * T7.13: parse_flash_info — too short data fails
 * ================================================================ */
static void
test_parse_flash_info_too_short (void)
{
  guint8 data[10];  /* less than 14 bytes */

  memset (data, 0, sizeof (data));

  ValidityFlashIcParams ic;
  guint16 num_partitions;

  gboolean ok = validity_pair_parse_flash_info (data, sizeof (data),
                                                &ic, &num_partitions);
  g_assert_false (ok);
}

/* ================================================================
 * T7.14: serialize_partition — known output format
 * ================================================================ */
static void
test_serialize_partition (void)
{
  ValidityPartition part = {
    .id = 1,
    .type = 3,
    .access_lvl = 0x0002,
    .offset = 0x1000,
    .size = 0x8000,
  };

  guint8 out[VALIDITY_PARTITION_ENTRY_SIZE];

  validity_pair_serialize_partition (&part, out);

  /* Check first 12 bytes: id(1) type(1) access_lvl(2LE) offset(4LE) size(4LE) */
  g_assert_cmpuint (out[0], ==, 1);
  g_assert_cmpuint (out[1], ==, 3);
  g_assert_cmpuint (FP_READ_UINT16_LE (out + 2), ==, 0x0002);
  g_assert_cmpuint (FP_READ_UINT32_LE (out + 4), ==, 0x1000);
  g_assert_cmpuint (FP_READ_UINT32_LE (out + 8), ==, 0x8000);

  /* Bytes 12-15 should be zero */
  for (int i = 12; i < 16; i++)
    g_assert_cmpuint (out[i], ==, 0);

  /* Bytes 16-47 = SHA-256 of the 12-byte entry */
  g_autoptr(GChecksum) checksum = g_checksum_new (G_CHECKSUM_SHA256);
  g_checksum_update (checksum, out, 12);
  guint8 expected_hash[32];
  gsize hash_len = 32;
  g_checksum_get_digest (checksum, expected_hash, &hash_len);

  g_assert_cmpmem (out + 16, 32, expected_hash, 32);

  /* Total size must be VALIDITY_PARTITION_ENTRY_SIZE */
  g_assert_cmpuint (sizeof (out), ==, VALIDITY_PARTITION_ENTRY_SIZE);
}

/* ================================================================
 * T7.15: make_cert — produces 444-byte certificate
 * ================================================================ */
static void
test_make_cert_size (void)
{
  guint8 pub_x[32], pub_y[32];

  RAND_bytes (pub_x, 32);
  RAND_bytes (pub_y, 32);

  gsize cert_len;
  g_autofree guint8 *cert = validity_pair_make_cert (pub_x, pub_y, &cert_len);

  g_assert_nonnull (cert);
  g_assert_cmpuint (cert_len, ==, VALIDITY_CLIENT_CERT_SIZE);

  /* First 4 bytes should be 0x17 in LE */
  g_assert_cmpuint (FP_READ_UINT32_LE (cert), ==, 0x17);
  /* Bytes 4-7 should be 0x20 in LE */
  g_assert_cmpuint (FP_READ_UINT32_LE (cert + 4), ==, 0x20);
}

/* ================================================================
 * T7.16: make_cert — deterministic for same input
 * ================================================================ */
static void
test_make_cert_deterministic (void)
{
  /* Using fixed keys so signature is reproducible.
   * Note: ECDSA uses random k, so signatures differ — but the
   * structure should be consistent. Actually ECDSA is NOT deterministic
   * without RFC 6979, so we can only verify structure matches. */
  guint8 pub_x[32] = { 0x01 };
  guint8 pub_y[32] = { 0x02 };

  gsize len1, len2;
  g_autofree guint8 *cert1 = validity_pair_make_cert (pub_x, pub_y, &len1);
  g_autofree guint8 *cert2 = validity_pair_make_cert (pub_x, pub_y, &len2);

  g_assert_nonnull (cert1);
  g_assert_nonnull (cert2);
  g_assert_cmpuint (len1, ==, len2);
  g_assert_cmpuint (len1, ==, VALIDITY_CLIENT_CERT_SIZE);

  /* Header and public key portion should be identical (first 184 bytes = cert body) */
  g_assert_cmpmem (cert1, 184, cert2, 184);
}

/* ================================================================
 * T7.17: encrypt_key — output blob has correct structure
 * ================================================================ */
static void
test_encrypt_key_structure (void)
{
  guint8 priv[32], pub_x[32], pub_y[32];
  guint8 enc_key[32], val_key[32];

  RAND_bytes (priv, 32);
  RAND_bytes (pub_x, 32);
  RAND_bytes (pub_y, 32);
  RAND_bytes (enc_key, 32);
  RAND_bytes (val_key, 32);

  gsize blob_len;
  g_autofree guint8 *blob = validity_pair_encrypt_key (priv, pub_x, pub_y,
                                                       enc_key, val_key,
                                                       &blob_len);

  g_assert_nonnull (blob);
  /* Blob format: 0x02(1) + IV(16) + ciphertext(112) + HMAC(32) = 161 */
  g_assert_cmpuint (blob_len, ==, 161);
  g_assert_cmpuint (blob[0], ==, VALIDITY_ENCRYPTED_KEY_PREFIX);
}

/* ================================================================
 * T7.18: encrypt_key — HMAC verification
 * ================================================================ */
static void
test_encrypt_key_hmac_valid (void)
{
  guint8 priv[32], pub_x[32], pub_y[32];
  guint8 enc_key[32], val_key[32];

  RAND_bytes (priv, 32);
  RAND_bytes (pub_x, 32);
  RAND_bytes (pub_y, 32);
  RAND_bytes (enc_key, 32);
  RAND_bytes (val_key, 32);

  gsize blob_len;
  g_autofree guint8 *blob = validity_pair_encrypt_key (priv, pub_x, pub_y,
                                                       enc_key, val_key,
                                                       &blob_len);

  g_assert_nonnull (blob);
  g_assert_cmpuint (blob_len, ==, 161);

  /* Blob: [0x02][iv:16][ct:112][hmac:32]
   * HMAC is over iv+ct (bytes 1..128) */
  const guint8 *iv_ct = blob + 1;
  gsize iv_ct_len = 16 + 112;
  const guint8 *stored_hmac = blob + 1 + iv_ct_len;

  guint8 computed_hmac[32];
  guint hmac_len = 32;
  HMAC (EVP_sha256 (), val_key, 32, iv_ct, iv_ct_len,
        computed_hmac, &hmac_len);

  g_assert_cmpmem (stored_hmac, 32, computed_hmac, 32);
}

/* ================================================================
 * T7.19: build_partition_flash_cmd — valid output structure
 * ================================================================ */
static void
test_build_partition_flash_cmd (void)
{
  /* Use a real device descriptor for the flash layout */
  const ValidityDeviceDesc *desc =
    validity_hal_device_lookup (VALIDITY_DEV_9A);

  g_assert_nonnull (desc);

  ValidityFlashIcParams flash_ic = {
    .size = 0x200000,
    .sector_size = 0x1000,
    .sector_erase_cmd = 0x20,
  };

  guint8 pub_x[32], pub_y[32];
  RAND_bytes (pub_x, 32);
  RAND_bytes (pub_y, 32);

  gsize cmd_len;
  g_autofree guint8 *cmd =
    validity_pair_build_partition_flash_cmd (&flash_ic, desc->flash_layout,
                                             pub_x, pub_y, &cmd_len);

  g_assert_nonnull (cmd);
  g_assert_cmpuint (cmd_len, >, 5);

  /* Command prefix: 0x4f followed by 4 zero bytes */
  g_assert_cmpuint (cmd[0], ==, 0x4f);
  g_assert_cmpuint (cmd[1], ==, 0);
  g_assert_cmpuint (cmd[2], ==, 0);
  g_assert_cmpuint (cmd[3], ==, 0);
  g_assert_cmpuint (cmd[4], ==, 0);
}

/* ================================================================
 * T7.20: build_tls_flash — produces exactly 4096 bytes
 * ================================================================ */
static void
test_build_tls_flash_size (void)
{
  ValidityPairState state;

  validity_pair_state_init (&state);

  /* Set up minimal test data */
  guint8 priv_blob[100];
  guint8 server_cert[200];
  guint8 ecdh_blob[400];
  RAND_bytes (priv_blob, sizeof (priv_blob));
  RAND_bytes (server_cert, sizeof (server_cert));
  RAND_bytes (ecdh_blob, sizeof (ecdh_blob));

  state.priv_blob = priv_blob;
  state.priv_blob_len = sizeof (priv_blob);
  state.server_cert = server_cert;
  state.server_cert_len = sizeof (server_cert);
  state.ecdh_blob = ecdh_blob;
  state.ecdh_blob_len = sizeof (ecdh_blob);

  gsize flash_len;
  g_autofree guint8 *flash = validity_pair_build_tls_flash (&state, &flash_len);

  g_assert_nonnull (flash);
  g_assert_cmpuint (flash_len, ==, 0x1000);

  /* Verify padding bytes at end are 0xff */
  gboolean has_ff_padding = FALSE;
  for (gsize i = flash_len - 1; i > 0; i--)
    {
      if (flash[i] == 0xff)
        {
          has_ff_padding = TRUE;
          break;
        }
    }
  g_assert_true (has_ff_padding);

  /* Don't free embedded pointers since they're stack-allocated */
  state.priv_blob = NULL;
  state.server_cert = NULL;
  state.ecdh_blob = NULL;
  validity_pair_state_free (&state);
}

/* ================================================================
 * T7.21: build_tls_flash — block structure
 * ================================================================ */
static void
test_build_tls_flash_blocks (void)
{
  ValidityPairState state;

  validity_pair_state_init (&state);

  guint8 priv_blob[50];
  guint8 server_cert[100];
  guint8 ecdh_blob[400];
  memset (priv_blob, 0xAA, sizeof (priv_blob));
  memset (server_cert, 0xBB, sizeof (server_cert));
  memset (ecdh_blob, 0xCC, sizeof (ecdh_blob));

  state.priv_blob = priv_blob;
  state.priv_blob_len = sizeof (priv_blob);
  state.server_cert = server_cert;
  state.server_cert_len = sizeof (server_cert);
  state.ecdh_blob = ecdh_blob;
  state.ecdh_blob_len = sizeof (ecdh_blob);

  gsize flash_len;
  g_autofree guint8 *flash = validity_pair_build_tls_flash (&state, &flash_len);
  g_assert_nonnull (flash);

  /* Block 0 should be first: id=0, size=1 */
  g_assert_cmpuint (FP_READ_UINT16_LE (flash), ==, 0);      /* block id */
  g_assert_cmpuint (FP_READ_UINT16_LE (flash + 2), ==, 1);  /* size = 1 */
  /* Skip 32-byte hash at flash+4 and 1-byte body at flash+36 */

  /* Next block should be block 4 (priv_blob) at offset 37 */
  gsize offset = 4 + 32 + 1;  /* header(4) + hash(32) + body(1) */
  g_assert_cmpuint (FP_READ_UINT16_LE (flash + offset), ==, 4);  /* block id */
  g_assert_cmpuint (FP_READ_UINT16_LE (flash + offset + 2), ==,
                    sizeof (priv_blob));

  state.priv_blob = NULL;
  state.server_cert = NULL;
  state.ecdh_blob = NULL;
  validity_pair_state_free (&state);
}

/* ================================================================
 * T7.22: pair state init and free
 * ================================================================ */
static void
test_pair_state_lifecycle (void)
{
  ValidityPairState state;

  validity_pair_state_init (&state);

  g_assert_null (state.client_key);
  g_assert_null (state.server_cert);
  g_assert_null (state.ecdh_blob);
  g_assert_null (state.priv_blob);
  g_assert_cmpuint (state.num_partitions, ==, 0);
  g_assert_cmpuint (state.erase_step, ==, 0);

  /* Free should be safe on empty state */
  validity_pair_state_free (&state);
}

/* ================================================================
 * T7.23: pair state free with allocated resources
 * ================================================================ */
static void
test_pair_state_free_with_resources (void)
{
  ValidityPairState state;

  validity_pair_state_init (&state);

  state.server_cert = g_malloc (100);
  state.server_cert_len = 100;
  state.ecdh_blob = g_malloc (400);
  state.ecdh_blob_len = 400;
  state.priv_blob = g_malloc (161);
  state.priv_blob_len = 161;

  /* Generate a key to test EVP_PKEY_free path */
  EVP_PKEY_CTX *pctx = EVP_PKEY_CTX_new_id (EVP_PKEY_EC, NULL);
  EVP_PKEY_keygen_init (pctx);
  EVP_PKEY_CTX_set_ec_paramgen_curve_nid (pctx, NID_X9_62_prime256v1);
  EVP_PKEY_keygen (pctx, &state.client_key);
  EVP_PKEY_CTX_free (pctx);
  g_assert_nonnull (state.client_key);

  /* Free should release all resources without leak */
  validity_pair_state_free (&state);
}

/* ================================================================
 * T7.24: encrypt_key — different inputs produce different blobs
 * ================================================================ */
static void
test_encrypt_key_different_inputs (void)
{
  guint8 priv1[32], priv2[32], pub_x[32], pub_y[32];
  guint8 enc_key[32], val_key[32];

  RAND_bytes (priv1, 32);
  RAND_bytes (priv2, 32);
  RAND_bytes (pub_x, 32);
  RAND_bytes (pub_y, 32);
  RAND_bytes (enc_key, 32);
  RAND_bytes (val_key, 32);

  gsize len1, len2;
  g_autofree guint8 *blob1 = validity_pair_encrypt_key (priv1, pub_x, pub_y,
                                                        enc_key, val_key, &len1);
  g_autofree guint8 *blob2 = validity_pair_encrypt_key (priv2, pub_x, pub_y,
                                                        enc_key, val_key, &len2);

  g_assert_nonnull (blob1);
  g_assert_nonnull (blob2);
  g_assert_cmpuint (len1, ==, len2);

  /* Different private keys should produce different ciphertexts */
  g_assert_true (memcmp (blob1, blob2, len1) != 0);
}

int
main (int argc, char *argv[])
{
  g_test_init (&argc, &argv, NULL);

  g_test_add_func ("/validity/pair/parse-flash-info-valid",
                   test_parse_flash_info_valid);
  g_test_add_func ("/validity/pair/parse-flash-info-needs-pairing",
                   test_parse_flash_info_needs_pairing);
  g_test_add_func ("/validity/pair/parse-flash-info-too-short",
                   test_parse_flash_info_too_short);
  g_test_add_func ("/validity/pair/serialize-partition",
                   test_serialize_partition);
  g_test_add_func ("/validity/pair/make-cert-size",
                   test_make_cert_size);
  g_test_add_func ("/validity/pair/make-cert-deterministic",
                   test_make_cert_deterministic);
  g_test_add_func ("/validity/pair/encrypt-key-structure",
                   test_encrypt_key_structure);
  g_test_add_func ("/validity/pair/encrypt-key-hmac-valid",
                   test_encrypt_key_hmac_valid);
  g_test_add_func ("/validity/pair/build-partition-flash-cmd",
                   test_build_partition_flash_cmd);
  g_test_add_func ("/validity/pair/build-tls-flash-size",
                   test_build_tls_flash_size);
  g_test_add_func ("/validity/pair/build-tls-flash-blocks",
                   test_build_tls_flash_blocks);
  g_test_add_func ("/validity/pair/state-lifecycle",
                   test_pair_state_lifecycle);
  g_test_add_func ("/validity/pair/state-free-with-resources",
                   test_pair_state_free_with_resources);
  g_test_add_func ("/validity/pair/encrypt-key-different-inputs",
                   test_encrypt_key_different_inputs);

  return g_test_run ();
}