//Basic tests for the C wrappers for rsa and mpi.
//S.MG, 2018

#include "mpi.h"
#include "smg_rsa.h"

#include "c_wrappers.h"

void test_mpi_cmp() {
  int result;
  int i;
  char a[KEY_LENGTH_OCTETS];
  char b[KEY_LENGTH_OCTETS];

  //initialize mpis
  for (i=0;i<KEY_LENGTH_OCTETS;i++) {
    a[i] = i % 256;
    b[i] = a[i];
  }

  //cmp same value mpi
  result = mpi_cmp_octets(a, KEY_LENGTH_OCTETS, b, KEY_LENGTH_OCTETS);
  if (result == 0)
    printf("PASS: mpi_cmp_octets on a == b.\n");
  else {
    printf("FAIL: mpi_cmp_octets on a == b ");
    printf("returned %d instead of 0.\n", result);
  }

  //cmp a < b
  b[489] = 234;
  result = mpi_cmp_octets(a, KEY_LENGTH_OCTETS, b, KEY_LENGTH_OCTETS);
  if (result == -1)
    printf("PASS: mpi_cmp_octets on a < b.\n");
  else {
    printf("FAIL: mpi_cmp_octets on a < b ");
    printf("returned %d instead of -1.\n", result);
  }

  //cmp a > b
  a[240] = 241;
  result = mpi_cmp_octets(a, KEY_LENGTH_OCTETS, b, KEY_LENGTH_OCTETS);
  if (result == 1)
    printf("PASS: mpi_cmp_octets on a > b.\n");
  else {
    printf("FAIL: mpi_cmp_octets on a > b ");
    printf("returned %d instead of 1.\n", result);
  }
}

void test_gen_rsa_octets() {
  RSA_secret_key sk;
  RSA_public_key pk;
	int nlimbs = mpi_nlimb_hint_from_nbytes( KEY_LENGTH_OCTETS );
	int nlimbs_pq = mpi_nlimb_hint_from_nbytes( KEY_LENGTH_OCTETS / 2 );

  //allocate memory
	sk.n = mpi_alloc(nlimbs);
	sk.e = mpi_alloc(nlimbs);
	sk.d = mpi_alloc(nlimbs);
	sk.p = mpi_alloc(nlimbs_pq);
	sk.q = mpi_alloc(nlimbs_pq);
	sk.u = mpi_alloc(nlimbs_pq);

	pk.n = mpi_alloc(nlimbs);
	pk.e = mpi_alloc(nlimbs);

  //generate key pair
  int len_n = KEY_LENGTH_OCTETS;
  int len_e = len_n;
  int len_d = len_n;
  int len_p = KEY_LENGTH_OCTETS / 2;
  int len_q = len_p;
  int len_u = len_p;
  char n[KEY_LENGTH_OCTETS];
  char e[KEY_LENGTH_OCTETS];
  char d[KEY_LENGTH_OCTETS];
  char p[KEY_LENGTH_OCTETS / 2];
  char q[KEY_LENGTH_OCTETS / 2];
  char u[KEY_LENGTH_OCTETS / 2];
  gen_rsa_octets(n, &len_n,
                 e, &len_e,
                 d, &len_d,
                 p, &len_p,
                 q, &len_q,
                 u, &len_u);

  //check encryption/decr works
  mpi_set_buffer(sk.n, n, len_n, 0);
  mpi_set_buffer(sk.e, e, len_e, 0);
  mpi_set_buffer(sk.d, d, len_d, 0);
  mpi_set_buffer(sk.p, p, len_p, 0);
  mpi_set_buffer(sk.q, q, len_q, 0);
  mpi_set_buffer(sk.u, u, len_u, 0);

  mpi_set_buffer(pk.n, n, len_n, 0);
  mpi_set_buffer(pk.e, e, len_e, 0);

  MPI encr = mpi_alloc(0);
	MPI plain = mpi_alloc(0);
  MPI out = mpi_alloc(0);
	mpi_fromstr(plain, "0x\
5B6A8A0ACF4F4DB3F82EAC2D20255E4DF3E4B7C799603210766F26EF87C8980E737579\
EC08E6505A51D19654C26D806BAF1B62F9C032E0B13D02AF99F7313BFCFD68DA46836E\
CA529D7360948550F982C6476C054A97FD01635AB44BFBDBE2A90BE06F7984AC8534C3\
28097EF92F6E78CAE0CB97");
  public_rsa(encr, plain, &pk);
  secret_rsa(out, encr, &sk);

  if (mpi_cmp(out, plain) != 0)
    printf("FAIL: test_gen_rsa encr/decr failed.\n");
  else
    printf("PASS: test_gen_rsa encr/decr passed.\n");

  //tidy up
  mpi_free(sk.n);
  mpi_free(sk.e);
  mpi_free(sk.d);
  mpi_free(sk.p);
  mpi_free(sk.q);
  mpi_free(sk.u);

  mpi_free(pk.n);
  mpi_free(pk.e);

  mpi_free(plain);
  mpi_free(encr);
  mpi_free(out);
}

void test_rsa_octets() {
  int noctets = 512;
	RSA_public_key pk;
	pk.n = mpi_alloc(0);
	pk.e = mpi_alloc(0);

  RSA_secret_key sk;
  sk.n = mpi_alloc(0);
  sk.e = mpi_alloc(0);
  sk.d = mpi_alloc(0);
  sk.p = mpi_alloc(0);
  sk.q = mpi_alloc(0);
  sk.u = mpi_alloc(0);

//key pair previously generated with EuCrypt
  mpi_fromstr(sk.n, "0x\
C6579F8646180EED0DC1F02E0DDD2B43EABB3F702D79D9928E2CDA5E1D42DF5D9ED777\
3F80B1F8D9B0DB7D4D00F55647640D70768F63D3CED56A39C681D08D6191F318BB79DC\
969B470A7364D53335C8318EF35E39D5DF706AB6F2393C6DD2128C142DBAB1806EB35E\
26C908F0A48419313D2D0F33DD430655DBFEC722899EC21C238E8DB7003430BBC39BAD\
990F9887F6B03E1344F537EC97389B78DBC656718ACD7B0FDC13DD24534F417BC7A18F\
077A0C4227354CEA19670331B6CAA3DFC17BBA7E70C14510D9EB3B63F3014994EC87BD\
23E868C0AE6E9EC55027577F62C0280B2D7DD1135001844923E5455C4566E066B3FDE9\
68C6BC4DC672F229FCE366440403D7A4F4A8BFBA5679B7D0844BA1231277D13A77C9E2\
B5A1CB138C1B7AB5B4D4832448723A3DE70ED2E86D5FC5174F949A02DE8E404304BEB9\
5F9BF40F3AA3CA15622D2776294BE7E19233406FF563CB8C25A1CB5AADBC1899DA3F2A\
E38533931FE032EE3232C2CD4F219FADF95B91635C0762A476A4DE5013F4384093F0FB\
715028D97F93B2E6F057B99EE344D83ADF2686FD5C9C793928BEF3182E568C4339C36C\
744C8E9CA7D4B9A16AA039CBF6F38CC97B12D87644E94C9DBD6BC93A93A03ED61ECC58\
74586E3A310E958F858735E30019D345C62E5127B80652C8A970A14B31F03B3A157CD5");

  mpi_fromstr(sk.e, "0x\
F74D78E382FC19B064411C6C20E0FDB2985F843007A54C7D8400BB459468624126E7D1\
75F397E55C57AF25858EAE2D2952FB7998C119A6103606733EB5E1D27FCA1FACF14ADE\
94101D383D1B25DA511805569BC344EAD384EDBF3F3A541B34887FE199D99D7F62E6E9\
D516F88D6F5AD3E020DF04D402A02CC628A0064362FE8516CF7CD6040E9521407AB90E\
E6B5AFFF9EA9EBB16A7D3407CE81FD3844F519880556AB94AB349C1F3BBB6FDB4C4B37\
7FE4C091EBDC2C3A1BD3AA56382D8D80E7742B5C751008FD6ECDD2EC3B2E3B6C566F69\
8ED672000B403766DD63C3ACBDE16A14FB02E83A2EB6AA018BFC0020401E790DEE24E9");

  mpi_fromstr(sk.d, "0x\
698DA05DA25B230211EEF0CBA12083A1457B749A11937AC9993859F69A3BF38D575E51\
66AF2EC88D77F1DF04E68AEA358EACF7659FD4722A4F5A1C8BA7676DA97A9FBA754511\
52F8F68887D3451A9CCFFFE9EB80979786E37495B17687A6212F77FA616E4C0CD8A8EB\
7AEB88EA6CCABB7F3E854FB94B35394A09F95F0D6F997947E865CC0606F437C30FE8C4\
8D96FBF5E2F52807BC9E9ED7BBEB23D5C45EDDCD16FE2BF410A9A1E5EF879E71C0D41F\
AE270C0C5D442860103F8C3944E802F33DB38432F11F763A7AF593656108E4A98A44A8\
549913CE5DCEC1A6500F280E3190991B2B938561CFACD8BC5183AAC9A4914BFE52C3BE\
39BB83688E1DE52479107EF8E087DCDB409432FC954C6349407E81DDFB11AE92BABB32\
A31868597958C9C76E0B4156F380955F0E09C1F3B98BB4CDD59E1B5C7D8CC2AA7491B0\
D319D219CF459A527CE1AA2729DEC53269653BF0ED3E0253F4451168437E3B069E4835\
0CA4C3EC82134E87135624C768D1330B0D70C6E447FD9945BF06FCB91AA334C0FD8EEF\
1ADBC15928B3DB62077B537F7E9F468CC95CD5AAFEAE1F760A863B48D07B163F670E2E\
5B550BB3E960230BA9FDAED9903AE2E669A7F3C4D1F1E25B8E8EDB8CC6E6FD2164E66F\
4E64ED77BEF1EC9E6CEA5624FD84C0680248746DC1C8187145F3CD2411659DAEAD11D");

  mpi_fromstr(sk.p, "0x\
CDD6F7673A501FB24C44D56CA1D434F6CB3334E193E02F8E906241906BCB7412DD2159\
825B24C22002F373E647C2DA62A854F3841C00FD5985D03227CA9B54A69380BA9D63BE\
738BDF9E65C247E43E1220EEDD9281DCA78B32A4E1B786B7697ED0C3195D5AF2990881\
B11D6FC9EC9F940067B2DEA2A516FAA5F269C98F0B67628A6D2708515A4A58041AA17A\
93E4C4DD95C85BC38351DDA1DCF3DFD91C505B22383132649CF9F9233852C7207075BC\
F43C71038F043F1EC53E9787FB051B7927D020903233C16897B993C8089D8464451F08\
6E756CF20E46CE6ED4A6AC5C327A0AAFBECBAAFD177969E7C952C76A4F4E7C85BF7F63");

  mpi_fromstr(sk.q, "0x\
F6ACF0790A250802C8D45DAC549CDBEF7806D5877A5DF0069136A458FAC4F0B0858060\
A873DA6355A965A064A0BC1BBB874872CD7ED89674AD407533041E74BCA317EC73597D\
335115523F61A05071E5ED81EE2A05331F65D4DC7A25AD7938B124CF03F49154B6693F\
B0B598B33ABDEF85C599A57A9B7347EAFF82638E1CBC28FCDFFF1FF04A18C2DBF39383\
95C2F8D1782B43D3A25EF7633B5DDAC89EFD3BAA64D976425A0891E00B876E9DE9FE4B\
6492B0EA8DFC7C8DEEC61721356EC816295B1BD9CD9DA3E30D2D90DC9CB3987F4BE042\
104900E036F3044A016749EF910CCFB9F377A90849B4CCCF4471A74E67EF6C814C9467");

  mpi_fromstr(sk.u, "0x\
854B89ED10F52258D00D6B3FA7F1FD22752804668F51FF7806DB82E22CB8B3AA8448D9\
B8E9DB14D31A36AEC2BCFA89E341B7334D494E97ED8051244136192233332C4612D963\
E7B6AF2535FDB7FE97E28DDFEBDFB3E1AFC29D05DBDF37106A817D3AB1864C7F7F2479\
82897EDA6A92BED47D9C68305CD170C7301ACEB05F8A6382E73CC7614B2D8D758669B3\
A99AB64114809254B0BE21F40341A5B48B9B032603B14875B87EB5E16603FD16552E14\
6A0FC6964958DFC25AA9FFCCD1ED1F4DEAF9FBAA0D7357F5FF0803FEB9BA78E74AC6B3\
070F417CEC6CFC7A3CF1E305FC7B76B7ED71893999AF797B2EBDE41FE90F076CCEDBFB");

// copy the public key components
  pk.n = mpi_copy( sk.n );
  pk.e = mpi_copy( sk.e );

// some plain text message
	MPI plain = mpi_alloc(0);
	mpi_fromstr(plain, "0x\
5B6A8A0ACF4F4DB3F82EAC2D20255E4DF3E4B7C799603210766F26EF87C8980E737579\
EC08E6505A51D19654C26D806BAF1B62F9C032E0B13D02AF99F7313BFCFD68DA46836E\
CA529D7360948550F982C6476C054A97FD01635AB44BFBDBE2A90BE06F7984AC8534C3\
28097EF92F6E78CAE0CB97");

// expected encrypted MPI (via rsa.c directly)
  MPI encr = mpi_alloc(0);
  public_rsa( encr, plain, &pk);
  MPI decr = mpi_alloc(0);
  secret_rsa( decr, encr, &sk);
  if (mpi_cmp(decr, plain) != 0)
    printf("FAIL: decrypted != plain in test_rsa_octets (MPI call)\n");

//allocate the char arrays for _octets rsa
  int len_n = noctets;
  int len_e = len_n;
  int len_d = len_n;
  int len_p = noctets / 2;
  int len_q = len_p;
  int len_u = len_p;
  char n[noctets];
  char e[noctets];
  char d[noctets];
  char p[noctets / 2];
  char q[noctets / 2];
  char u[noctets / 2];

//copy the key components into char arrays
  mpi_to_octets(n, &len_n, sk.n);
  mpi_to_octets(e, &len_e, sk.e);
  mpi_to_octets(d, &len_d, sk.d);
  mpi_to_octets(p, &len_p, sk.p);
  mpi_to_octets(q, &len_q, sk.q);
  mpi_to_octets(u, &len_u, sk.u);

//call _octets rsa and check results
  int len_encr = noctets;
  int len_decr = noctets;
  int len_plain = noctets;
  char plain_o[noctets];
  char encr_o[noctets];
  char expected_encr_o[noctets];
  char decr_o[noctets];
  char expected_decr_o[noctets];

  mpi_to_octets(plain_o, &len_plain, plain);
  mpi_to_octets(expected_encr_o, &len_encr, encr);
  mpi_to_octets(expected_decr_o, &len_decr, decr);
  len_decr = noctets;

  int len;
  len = public_rsa_octets( encr_o, len_encr, plain_o, len_plain,
                           n, len_n, e, len_e);
  if (len != len_encr)
    printf("FAIL: actual len of encr is %d; expected %d\n", len, len_encr);
  else
    printf("PASS: actual len of encr matches expected: %d\n", len);
  int errors= 0;
  int i;
  for (i=0;i<len;i++)
    if (encr_o[i] != expected_encr_o[i])
      errors = errors + 1;
  if (errors>0)
    printf("FAIL: found %d errors in public_rsa_octets output\n", errors);
  else
    printf("PASS: no errors found in public_rsa_octets output\n");

  len_encr = len;
  len = private_rsa_octets( decr_o, len_decr, encr_o, len_encr,
                            n, len_n, e, len_e,
                            d, len_d, p, len_p,
                            q, len_q, u, len_u);
  if (len != len_plain)
    printf("FAIL: actual len of decr is %d; expected %d\n", len, len_plain);
  else
    printf("PASS: actual len of decr matches expected: %d\n", len);

  errors = 0;
  for (i=0;i<len;i++)
    if (decr_o[i] != expected_decr_o[i])
      errors = errors + 1;
  if (errors>0)
    printf("FAIL: found %d errors in private_rsa_octets output\n", errors);
  else printf("PASS: no errors found in private_rsa_octets_output\n");

  //tidy up
  mpi_free(sk.n);
  mpi_free(sk.e);
  mpi_free(sk.d);
  mpi_free(sk.p);
  mpi_free(sk.q);
  mpi_free(sk.u);

  mpi_free(pk.n);
  mpi_free(pk.e);

  mpi_free(plain);
  mpi_free(encr);
  mpi_free(decr);
}

int main(int ac, char **av) {
  if (ac < 2) {
    printf("Usage: %s testID\n", av[0]);
    return -1;
  }

  int id = atoi(av[1]);
  switch (id) {
    case 1:
      test_mpi_cmp();
      break;
    case 2:
      test_gen_rsa_octets();
      break;
    case 3:
      test_rsa_octets();
      break;
    default:
      printf("Current test ids:\n");
      printf("1 test of mpi_cmp_octets\n");
      printf("2 test of gen_rsa_octets (can be very SLOW!)\n");
      printf("3 test of rsa_octets (can take a few minutes)\n");
  }
  return 0;
}