//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\
CD2C025323BEA46FFF2FA8D7A9D39817EA713421F4AE03FA8120641193892A70BFECF5\
83101635A432110D3DDE6339E3CC7ECC0AD91C026FCACE832DD3888A6FCA7BCE56C390\
5A5AC8C7BC921DA675E4B62489B254EB34659D547D71165BC998983A81937BD251AEE1\
2D985EC387D5376F5DCC5EF7EC530FBD6FD2AA7285EE1AF3335EA73163F0954F30402E\
D7B374EE84A97B1849B0674B0DA0A2050BD79B71ABB1559F3A9CFDB8557DED7BC90CF2\
09E8A847E9C226140845B7D03842162E7DA5DD16326CB1F71A248D841FE9076A09911F\
2F4F5E3EA44EA8DE40332BF00406990BCCF61C322A03C456EF3A98B341E0BDBC1088CE\
683E78510E76B72C2BCC1EE9AEDD80FFF18ABFC5923B2F36B581C25114AB2DF9F6C2B1\
9481703FD19E313DCD7ACE15FA11B27D25BCE5388C180A7E21167FB87750599E1ED7C7\
50F4A844E1DC2270C62D19671CF8F4C25B81E366B09FC850AE642136D204A9160AEECE\
575B57378AA439E9DD46DC990288CD54BAA35EEE1C02456CD39458A6F1CBF012DCEDF4\
27CCF3F3F53645658FC49C9C9D7F2856DB571D92B967AB5845514E0054DDB49099F5DD\
04A6F6F5C5CE642276834B932881AEB648D1F25E9223971F56E249EF40CF7D80F22621\
CDD0260E9E7D23746960ADB52CF2987584FB1DE95A69A39E5CB12B76E0F5C1A0529C0C\
065D2E35720810F7C7983180B9A9EA0E00C11B79DC3D");

  mpi_fromstr(sk.e, "0x\
DD4856B4EE3D099A8604AE392D8EFEC094CDF01546A28BE87CB484F999E8E75CDFCD01\
D04D455A6A9254C60BD28C0B03611FC3E751CC27EF768C0B401C4FD2B27C092834A6F2\
49A145C4EDC47A3B3D363EC352462C945334D160AF9AA72202862912493AC6190AA3A6\
149D4D8B9996BA7927D3D0D2AD00D30FD630CF464E6CAF9CF49355B9A70E05DB7AE915\
F9F602772F8D11E5FCDFC7709210F248052615967090CC1F43D410C83724AA5912B2F0\
52E6B39449A89A97C79C92DC8CB8DEEFCF248C1E1D2FC5BFE85165ECA31839CAA9CEB3\
3A92EBDC0EB3BAC0F810938BB173C7DA21DCBB2220D44CBA0FD40A2C868FC93AC5243E\
C137C27B0A76D65634EBB3");

  mpi_fromstr(sk.d, "0x\
7C8A6FA1199D99DCA45E9BDF567CA49D02B237340D7E999150BC4883AE29DEC5158521\
B338F35DC883792356BDDBB3C8B3030A6DD4C6522599A3254E751F9BA1CB1061C5633C\
81BBFACF6FCD64502614102DFED3F3FA284066C342D5E00953B415915331E30812E5FB\
CD6680ADCCDEE40B8376A3A225F2E160EA59C7566804526D73BB660A648A3EF9802313\
B2F841E8458B2AAACE7AACF31083E8F3F630298138393BC88BBD7D4AA4334949651D25\
365B10DBF4A4A08E20A6CC74BFDD37C1C38E2ADC2A283DF06590DF06B46F67F6ACA67F\
AC464C795261659A2F9558802D0BBAA05FD1E1AF2CDC70654723DF7EFAEA148B8CDBEB\
C89EA2320AB9BBB1BC4311475DF3D91446F02EF192368DFEBAC598CCFD4407DEC58FDC\
1A94CCDD6E5FBA9C52164ACEA8AEE633E557BCCEACB7A1AF656C379482D784A120A725\
32F9B2B35173D505F21D5AD4CB9511BC836DC923730B70291B70290A216CA3B21CFF79\
E895C35F4F7AF80E1BD9ED2773BD26919A76E4298D169160593E0335BE2A2A2D2E8516\
948F657E1B1260E18808A9D463C108535FB60B3B28F711C81E5DE24F40214134A53CE5\
9A952C8970A1D771EBEFFA2F4359DCF157995B3F1950DE3C6EC41B7FF837148F55F323\
372AF3F20CE8B8038E750C23D8F5041FA951327859B0E47483F0A47103EF808C72C251\
006FA526245291C8C84C12D2EF63FB2301EA3EEDA42B");

  mpi_fromstr(sk.p, "0x\
E236732452039C14EC1D3B8095BDDCFB7625CE27B1EA5394CF4ED09D3CEECAA4FC0BF6\
2F7CE975E0C8929CE84B0259D773EA038396479BF15DA065BA70E549B248D77B4B23ED\
A267308510DBEE2FD44E35D880EE7CFB81E0646AA8630165BD8988C3A8776D9E704C20\
AA25CA0A3C32F27F592D5FD363B04DD57D8C61FFDCDFCCC59E2913DE0EE47769180340\
E1EA5A803AA2301A010FF553A380F002601F0853FCACDB82D76FE2FACBCD6E5F294439\
0799EA5AE9D7880D4E1D4AE146DC1D4E8495B9DD30E57E883923C5FC26682B7142D35C\
D8A0FC561FE725A6CF419B15341F40FE0C31132CBD81DD8E50697BD1EBFFA16B522E16\
F5B49A03B707218C7DA60B");

  mpi_fromstr(sk.q, "0x\
E830482A3C4F5C3A7E59C10FF8BA760DB1C6D55880B796FFDA4A82E0B60E974E81D04B\
2A4AD417823EBFB4E8EFB13782943562B19B6C4A680E3BA0C8E37B5023470F4F1AC1F8\
A0B10672EF75CD58BCD45E6B14503B8A6A70AFE79F6201AF56E7364A1C742BE1453FD2\
24FDC9D66522EAF4466A084BCB9E46D455A2946E94CBF028770F38D0B741C2CC59308F\
71D8C2B4B9C928E0AE8D68DEB48A3E9EFD84A10301EBD55F8221CA32FC567B306B2A8E\
116350AFB995859FDF4378C5CFD06901494E8CFA5D8FAC564D6531FA8A2E4761F5EFBA\
F78750B6F4662BE9EA4C2FAD67AF73EEB36B41FC15CB678810C19A51DF23555695C4C1\
546F3FACA39CAA7BB8DBD7");

  mpi_fromstr(sk.u, "0x\
846232322775C1CD7D5569DC59E2F3E61A885AE2E9C4A4F8CB3ACBE8C3A5441E5FE348\
A2A8AC9C2998FBF282222BF508AA1ECF66A76AEDD2D9C97028BFD3F6CA0542E38A5312\
603C70B95650CE73F80FDD729988FBDB5595A5BF8A007EA34E54994A697906CE56354C\
E00DF10EB711DEC274A62494E3D350D88736CF67A477FB600AC9F1D6580727585092BF\
5EBC092CC4D6CF75769051033A1197103BE269942F372168A53771746FBA18ED6972D5\
0B935A9B1D6B5B3DD50CD89A27FE93C10924E9103FACF7B4C5724A046C3D3B50CC1C78\
5F5C8E00DBE1D6561F120F5294C170914BC10F978ED4356EED67A9F3A60D70AFE540FC\
5373CBAE3D0A7FD1C87273");

// 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;
}