--S.MG, 2018

with Interfaces; use Interfaces;
with Interfaces.C; use Interfaces.C;
with Ada.Text_IO; use Ada.Text_IO;
with Ada.Sequential_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;

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 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
    ReadRSAKey( "keys_rsa.txt", 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...");
    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
    ReadRSAKey( "keys_rsa.txt", 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 Hex2Octets( Hex: in String; O: out Raw_Types.Octets ) is
    S : String := "16#AA#";
    -- to make sure that input String has EVEN number of chars (ie full octets)
    H : String(1..Hex'Length+Hex'Length mod 2) := (others=>'0');
  begin
    -- first char is 0 if needed to cover full octet...
    H(H'Length-Hex'Length+1..H'Length) := Hex;
    O := (others => 0);
    for I in 0 .. H'Length/2-1 loop
      S := "16#" & H(H'First + I*2 .. H'First + I*2 + 1) & "#";
      O(O'Last - H'Length/2 + 1 + I) := Unsigned_8'Value(S);
    end loop;
  end Hex2Octets;

  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;

  procedure ReadRSAKey( Filename: in String; Key: out RSA_OAEP.RSA_skey ) is
    package Char_IO is new Ada.Sequential_IO(Character);
    use Char_IO;
    Full : String(1..RSA_len'Length*2) := (others => '0');
    Half : String(1..RSA_half'Length*2) := (others => '0');
    F    : Char_IO.File_Type;
    C    : Character;
  begin
    Open( File => F, Mode => In_File, Name => Filename );

    -- read n
    for I in Full'Range loop
      Read(F, Full(I));
    end loop;
    -- read new line character and convert to hex
    Read(F, C);
    Hex2Octets(Full, Key.n);

    -- read e
    for I in Half'Range loop
      Read(F, Half(I));
    end loop;
    -- read new line character and convert to hex
    Read(F, C);
    Hex2Octets(Half, Key.e);

    -- read d
    for I in Full'Range loop
      Read(F, Full(I));
    end loop;
    -- read new line character and convert to hex
    Read(F, C);
    Hex2Octets(Full, Key.d);

    -- read p
    for I in Half'Range loop
      Read(F, Half(I));
    end loop;
    -- read new line character and convert to hex
    Read(F, C);
    Hex2Octets(Half, Key.p);

    -- read q
    for I in Half'Range loop
      Read(F, Half(I));
    end loop;
    -- read new line character and convert to hex
    Read(F, C);
    Hex2Octets(Half, Key.q);

    -- read u
    for I in Half'Range loop
      Read(F, Half(I));
    end loop;
    Hex2Octets(Half, Key.u);

    -- Close file
    Close( F );

    exception 
      when Char_IO.End_Error =>
        Put_Line("ReadRSAKey ERROR: Unexpected end of file in " & Filename);
      when others =>
        Put_Line("ReadRSAKey ERROR: can not open file " & Filename);

  end ReadRSAKey;

end Test_RSA_OAEP;