--S.MG, 2018

with Interfaces; use Interfaces;
with Interfaces.C; use Interfaces.C;
with Ada.Text_IO; use Ada.Text_IO;
with RSA_OAEP; use RSA_OAEP;
with OAEP; use OAEP;
with Raw_Types; use Raw_Types;
with RNG; use RNG;
with Keccak; use Keccak;
with IO_RSA;

package body Test_RSA_OAEP is

  procedure test_char_array is
    S : String := OAEP.TMSR_STR;
    O : Octets := Raw_Types.OAEP_RESERVED;
    A : char_array(0..O'Length-1) := (others => '0');
    B : Octets(0..O'Length -1) := (others => 0);
    Fail : Boolean := FALSE;
  begin
    Octets_To_Char_Array(O, A);
    Char_Array_To_Octets(A, B);

    if B /= O then
      Put_Line("FAIL: char_array_to_octets");
    else
      Put_Line("PASS: char_array_to_octets");    
    end if;

    for I in 0..S'Length-1 loop 
      declare
        C : Character := Character(A(A'First + size_t(I)));
        E : Character := S(S'First + I);
      begin
        if C /= E then
          Fail := TRUE;
          Put("Error at pos " & Integer'Image(I) & ": "); 
          Put(Integer'Image(Character'Pos(C)));
          Put_Line(" instead of " & Integer'Image(Character'Pos(E)));
        end if;
      end;
    end loop;
    if FAIL then
      Put_Line("FAIL: test octets_to_char_array");
    else
      Put_Line("PASS: test octets_to_char_array");
    end if;
  end test_char_array;

  -- test OAEP encrypt + decrypt
  procedure test_oaep is
    Plain: Octets(1..MAX_LEN_MSG);
    Short: Octets(0..10);
    Encr : OAEP_Block;
    Decr : OAEP_HALF;
    Len  : Natural;
    Entropy: OAEP_Block;
    Success : Boolean;
  begin
    RNG.Get_Octets(Plain);
    RNG.Get_Octets(Entropy);
    RNG.Get_Octets(Short);

    -- test full length message
    OAEP_Encrypt(Plain, Entropy, Encr);
    OAEP_Decrypt(Encr, Len, Decr, Success);

    if not Success or Len/8 /= Plain'Length then
      Put_Line("FAIL: oaep encrypt/decrypt on max len message.");
    else 
      if Decr(Decr'First..Decr'First+Len/8-1) /= 
              Plain(Plain'First..Plain'First+Len/8-1) then
        Put_Line("FAIL: oaep encrypt/decrypt on max len message - " &
                 "result different from expected.");    
      else
        Put_Line("PASS: oaep encrypt/decrypt on max len message.");
      end if;
    end if;

    -- test short message
    OAEP_Encrypt(Short, Entropy, Encr);
    OAEP_Decrypt(Encr, Len, Decr, Success);
    if not Success or Len/8 /= Short'Length then
      Put_Line("FAIL: oaep encrypt/decrypt on short message.");
    else 
      if Decr(Decr'First..Decr'First+Len/8-1) /= 
             Short(Short'First..Short'First+Len/8-1) then
        Put_Line("FAIL: oaep encrypt/decrypt on short message - " &
                 "result different from expected.");    
      else
        Put_Line("PASS: oaep encrypt/decrypt on short message.");
      end if;
    end if;
    
  end test_oaep;

  -- test JUST RSA (i.e. without oaep) with RSA key pair previously generated
  procedure test_rsa( E_Len    : in Positive; 
                      D_Len    : in Positive;
                      Filename : in String) is 
    Plain: OAEP_Block := (others => 0);
    Decr : OAEP_Block := (others => 0);
    Encr : RSA_len;
    pkey: RSA_pkey;
    skey: RSA_skey;
  begin
    -- initialize with RSA pair previously generated
    IO_RSA.ReadRSAKey( Filename, E_Len, D_Len, skey );

    -- copy n and e for public key
    pkey.n := skey.n;
    pkey.e := skey.e;
    -- get random data
    RNG.Get_Octets(Plain);
    -- make first octet < RSA key's modulus first octet
    Plain(Plain'First) := 16#00#;
    -- naked rsa encrypt/decrypt 
    Put_Line("Encrypting with RSA public key with e len " & 
         Positive'Image(E_Len) & "...");
    Public_RSA( Plain, pkey, Encr );
    Put_Line("Decrypting with RSA private key...");
    Private_RSA( Encr, skey, Decr );
    Put_Line("Checking...");

    -- check result
    if Decr /= Plain then
      Put_Line("FAIL: RSA encrypt/decrypt result doesn't match plain.");
    else
      Put_Line("PASS: RSA encrypt/decrypt");
    end if;
  end test_rsa;

  -- test rsa+oaep with RSA key pair previously generated
  procedure test_rsa_oaep is
    Plain: Octets(1..MAX_LEN_MSG) := (others=>20);
    Short: Octets(1..10);
    Decr : RSA_len;
    Encr : RSA_len;
    pkey: RSA_pkey;
    skey: RSA_skey;
    Success: Boolean;
    Len : Natural;
  begin
    -- initialize with RSA pair previously generated
    IO_RSA.ReadRSAKey( "keys_rsa.txt", Raw_Types.RSA_KEY_OCTETS, 
                       Raw_Types.RSA_KEY_OCTETS*2, skey );
    -- copy n and e for public key
    pkey.n := skey.n;
    pkey.e := skey.e;

    -- test with 0 message of length Plain'Length
    RSA_OAEP.Encrypt(Plain, pkey, Encr);
    RSA_OAEP.Decrypt(Encr, skey, Decr, Len, Success);
    if (not Success) or Len /= Plain'Length 
         or Plain /= Decr(Decr'First..Decr'First+Plain'Length-1) then
      Put_Line("FAIL: RSA_OAEP on max len message 20-filled.");
    else
      Put_Line("PASS: RSA_OAEP on max len message 20-filled.");
    end if;

    -- get random data for "plain" message
    RNG.Get_Octets(Plain);
    RSA_OAEP.Encrypt(Plain, pkey, Encr);
    RSA_OAEP.Decrypt(Encr, skey, Decr, Len, Success);
    if (not Success) or Len /= Plain'Length 
         or Plain /= Decr(Decr'First..Decr'First+Plain'Length-1) then
      Put_Line("FAIL: RSA_OAEP on random data of max length.");
    else
      Put_Line("PASS: RSA_OAEP on random data of max length.");
    end if;

    -- get random data for "short" message
    RNG.Get_Octets(Short);
    RSA_OAEP.Encrypt(Short, pkey, Encr);
    RSA_OAEP.Decrypt(Encr, skey, Decr, Len, Success);
    if (not Success) or Len /= Short'Length 
         or Short /= Decr(Decr'First..Decr'First+Short'Length-1) then
      Put_Line("FAIL: RSA_OAEP on random data of short length.");
    else
      Put_Line("PASS: RSA_OAEP on random data of short length.");
    end if; 

  end test_rsa_oaep;

-- helper methods

  procedure PrintOctets( O: in Raw_Types.Octets; Title: in String ) is
  begin
    Put_Line(Title);
    for V of O loop
      Put(Unsigned_8'Image(V) & " ");
    end loop;
    New_Line;
  end PrintOctets;

end Test_RSA_OAEP;