#include <assert.h>
#include <string.h>

#include <nettle/eddsa.h>

#include "test-utils.h"

#include "checkpoint.h"
#include "crypto-internal.h"
#include "policy.h"
#include "quorum.h"

#define X_P(x) ((BYTECODE_OP_X_P << 6) | (x))
#define GE_K(k) ((BYTECODE_OP_GE_K << 6) | (k))

static void
make_key_pair (struct sigsum_key_with_hash *pub, uint8_t priv[ED25519_KEY_SIZE],
	       unsigned id)
{
  memset (priv, id, ED25519_KEY_SIZE);
  ed25519_sha512_public_key (pub->key.k, priv);
  sigsum_hash_key (&pub->hash, &pub->key);
}

/* Sort witnesses and cosignatures by key hash. */
static void
sort_witnesses (size_t n, struct sigsum_key_with_hash *keys,
		struct sigsum_cosignature *cosignatures)
{
  /* Do a simple selection sort. */
  for (; n > 1; n--, keys++, cosignatures++)
    {
      size_t i, j;
      /* Find the smallest entry */
      for (i = 0, j = 1; j < n; j++)
	{
	  int c = memcmp (keys[i].hash.h, keys[j].hash.h, sizeof (keys[i].hash.h));
	  assert (c != 0);
	  if (c > 0)
	    i = j;
	}
      if (i > 0)
	{
	  /* Swap, to get smallest at index 0. */
	  struct sigsum_key_with_hash tmp_key = keys[0];
	  keys[0] = keys[i];
	  keys[i] = tmp_key;
	  struct sigsum_cosignature tmp_cs = cosignatures[0];
	  cosignatures[0] = cosignatures[i];
	  cosignatures[i] = tmp_cs;
	}
    }
}

static void
sign_tree_head (const struct sigsum_key_with_hash *pub, const uint8_t priv[ED25519_KEY_SIZE],
		const struct sigsum_tree_head *th, struct sigsum_signature *sig)
{
  uint8_t checkpoint[CHECKPOINT_BODY_MAX_SIZE];
  size_t checkpoint_size = _sigsum_checkpoint_body (checkpoint, &pub->hash, th);
  ed25519_sha512_sign (pub->key.k, priv, checkpoint_size, checkpoint, sig->s);
}

static void
cosign_tree_head (const struct sigsum_key_with_hash *pub, const uint8_t priv[ED25519_KEY_SIZE],
		  const struct sigsum_hash *log_kh, const struct sigsum_tree_head *th,
		  uint64_t timestamp, struct sigsum_cosignature *cs)
{
  uint8_t data[COSIGNATURE_PREFIX_MAX_SIZE + CHECKPOINT_BODY_MAX_SIZE];
  uint8_t *checkpoint = data + COSIGNATURE_PREFIX_MAX_SIZE;
  size_t checkpoint_size = _sigsum_checkpoint_body (checkpoint, log_kh, th);
  size_t prefix_size = _sigsum_cosignature_prefix (data, timestamp);
  ed25519_sha512_sign (pub->key.k, priv, checkpoint_size + prefix_size, checkpoint - prefix_size, cs->signature.s);
  cs->key_hash = pub->hash;
  cs->timestamp = timestamp;
}

static void
test_log_policy (void)
{
  static const struct sigsum_tree_head th = {10, {{1, 2, 3}}};
  struct sigsum_key_with_hash log_keys[3];
  struct sigsum_signature sigs[3];
  for (unsigned i = 0; i < 3; i++)
    {
      uint8_t priv[ED25519_KEY_SIZE];
      make_key_pair (&log_keys[i], priv, i+1);
      sign_tree_head (&log_keys[i], priv, &th, &sigs[i]);
    }
  static const uint8_t quorum[] = { BYTECODE_INSTR_TRUE };
  bool scratch[1];
  struct sigsum_policy policy;
  policy.num_logs = 2;
  policy.num_witnesses = 0;
  policy.logs = log_keys;
  policy.witnesses = NULL;
  policy.quorum_length = sizeof(quorum);
  policy.quorum_bytecode = quorum;
  policy.quorum_stack_size = 1;

  assert (sigsum_policy_verify_tree_head_itch (&policy) <= sizeof (scratch));

  const char *err;
  if (!sigsum_policy_verify_tree_head (&policy, &th, &log_keys[0].hash,
				       &sigs[0], 0, NULL,
				       &err, scratch))
    die ("verifying treehead for log 0 failed: %s", err);

  if (!sigsum_policy_verify_tree_head (&policy, &th, &log_keys[1].hash,
				       &sigs[1], 0, NULL,
				       &err, scratch))
    die ("verifying treehead for log 1 failed: %s", err);

  if (sigsum_policy_verify_tree_head (&policy, &th, &log_keys[2].hash,
				      &sigs[2], 0, NULL,
				      &err, scratch))
    die ("verifying treehead for log 2 succeeded, but it's not allowed by policy");

  if (sigsum_policy_verify_tree_head (&policy, &th, &log_keys[1].hash,
				      &sigs[0], 0, NULL,
				      &err, scratch))
    die ("verifying treehead for log 0 with log hash 1 succeeded");
}

static void
test_witness_policy (void)
{
  static const struct sigsum_tree_head th = {10, {{1, 2, 3}}};
  struct sigsum_key_with_hash log_key;
  struct sigsum_signature log_sig;
  uint8_t priv[ED25519_KEY_SIZE];
  make_key_pair (&log_key, priv, 1);
  sign_tree_head (&log_key, priv, &th, &log_sig);
#define NUM_WITNESSES 5
  struct sigsum_key_with_hash witness_keys[NUM_WITNESSES];
  struct sigsum_cosignature cosignatures[NUM_WITNESSES];
  for (unsigned i = 0; i < NUM_WITNESSES; i++)
    {
      make_key_pair (&witness_keys[i], priv, 10 + i);
      cosign_tree_head (&witness_keys[i], priv, &log_key.hash, &th, 1000 + i, &cosignatures[i]);
    }
  sort_witnesses (NUM_WITNESSES, witness_keys, cosignatures);

  /* 3-of-4 policy. */
  static uint8_t quorum[] = {
    X_P(0), X_P(1), BYTECODE_INSTR_ADD,
    X_P(2), BYTECODE_INSTR_ADD, X_P(3), BYTECODE_INSTR_ADD,
    GE_K (3)
  };

  struct sigsum_policy policy;
  policy.num_logs = 1;
  policy.num_witnesses = 4;
  policy.logs = &log_key;
  policy.witnesses = witness_keys;
  policy.quorum_length = sizeof(quorum);
  policy.quorum_bytecode = quorum;
  policy.quorum_stack_size = 2;

  bool scratch[6];
  assert (sigsum_policy_verify_tree_head_itch (&policy) <= sizeof (scratch));

#define LIST(n, ...) n, (LITERAL_STATIC const int[]){__VA_ARGS__}
  LITERAL_STATIC const struct {
    const char *desc;
    unsigned num_witnesses;
    const int *witnesses; /* Indices of witnesses to include. */
    int invalidate; /* Signature to invalidate (-1 if none) */
    bool expect_valid;
  } tests[] = {
    {"no cosignature", 0, NULL, -1, false},
    {"only one cosignature", LIST(1, 0), -1, false},
    {"only two cosignatures",LIST(2, 1, 4), -1, false},
    {"three cosignature", LIST(3, 0, 1, 2), -1, true},
    {"other three cosignature", LIST(3, 1, 2, 3), -1, true},
    {"all cosignatures", LIST(5, 0, 1, 2, 3, 4), -1, true},
    {"all cosignatures, one invalid", LIST(5, 0, 1, 2, 3, 4), 2, true},
    {"three cosignatures, but one invalid", LIST(3, 2, 3, 4), 2, false},
  };

  for (unsigned i = 0; i < sizeof (tests) / sizeof (tests[0]); i++)
    {
      struct sigsum_cosignature present[NUM_WITNESSES];
      for (unsigned j = 0; j < tests[i].num_witnesses; j++)
	{
	  present[j] = cosignatures[tests[i].witnesses[j]];
	  if (tests[i].witnesses[j] == tests[i].invalidate)
	    present[j].timestamp ^= 1;
	}
      const char *err;
      bool valid = sigsum_policy_verify_tree_head (&policy, &th, &log_key.hash,
						   &log_sig, tests[i].num_witnesses, present,
						   &err, scratch);
      if (tests[i].expect_valid && !valid)
	die ("%s: Failed on valid th: %s", tests[i].desc, err);
      if (!tests[i].expect_valid && valid)
	die ("%s: Expected error, but got none", tests[i].desc);
    }
}

int
main (void)
{
  test_log_policy ();
  test_witness_policy ();

  return 0;
}
