-- S.MG, 2018

with Interfaces; use Interfaces;
with OAEP;
with RNG;

package body RSA_OAEP is

  -- OAEP + RSA with given public key.
  -- Steps:
  -- 1. repeat oaep.encrypt of Plain 
  --      until the result has first octet < than first octet of Key's modulus.
  -- 2. RSA.encrypt with Key on the result from 1.
  procedure Encrypt( Plain: in Raw_Types.Octets;
                     Key  : in RSA_pkey;
                     Encr : out Raw_Types.RSA_len ) is
    -- block obtained through OAEP padding of given input Plain
    Blk     : OAEP.OAEP_Block := ( others => 0 );
    Entropy : OAEP.OAEP_Block;
    Len     : constant Natural := entropy'Length;
  begin
    -- loop cond: 1st octet of oaep block is < 1st octet of key's modulus
    loop
      -- get random bits via RNG (FG)
      RNG.Get_Octets( Entropy );

      -- oaep encrypt
      OAEP.OAEP_Encrypt( Plain, Entropy, Blk );

      -- check that the oaep block is suitable i.e. comparison of 1st octet
      if Blk(Blk'First) < Key.n(Key.n'First) then exit; end if;

    end loop; -- oaep block < modulus

    -- RSA encrypt
    Public_RSA( Blk, Key, Encr );
  end; --Encrypt (OAEP+RSA)

  -- RSA+OAEP Decrypt
  -- Steps:
  -- 1. RSA Decrypt (secret key)
  -- 2. OAEP Decrypt
  procedure Decrypt( Encr : in Raw_Types.RSA_len;
                     Key  : in RSA_skey;
                     Plain: out Raw_Types.Octets;
                     Plain_Len: out Natural;
                     Success: out Boolean) is
    Blk     : OAEP.OAEP_Block;
    Msg     : OAEP.OAEP_HALF;
    Msg_Len : Natural;
  begin
    -- RSA decrypt with provided secret key
    -- NB: result HAS TO BE AN OAEP BLOCK HERE! Even if - potentially - 0 led!
    Private_RSA( Encr, Key, Blk );

    -- OAEP decrypt
    OAEP.OAEP_Decrypt(Blk, Msg_Len, Msg, Success);
    -- switch to Length in OCTETS!
    Msg_Len := Msg_Len / 8;

    -- check that result FITS in given output array - otherwise, fail
    if Msg_Len > Plain'Length then
      Success := FALSE;
      Plain_Len := Msg_Len; --to give a clue to caller
    else
      Plain_Len := Msg_Len;
      Plain( Plain'First .. 
             Plain'First + Plain_Len-1) := Msg( Msg'First ..
                                                Msg'First + Plain_Len -1 );
    end if;
  end Decrypt;

  -- helper methods
  -- encrypt with public RSA key given as struct, Ada style
  procedure Public_RSA( Plain: in Raw_Types.Octets;
                        Key  : in RSA_pkey;
                        Encr : out Raw_Types.RSA_len) is
    Encr_char  : char_array( size_t(Encr'First) .. size_t(Encr'Last));
    Plain_char : char_array( size_t(Plain'First) .. size_t(Plain'Last));
    N_char     : char_array( size_t(Key.n'First) .. size_t(Key.n'Last));
    E_char     : char_array( size_t(Key.e'First) .. size_t(Key.e'Last));
    out_len    : Integer;
  begin
    -- convert to char array
    Octets_To_Char_Array( Plain, Plain_char );
    Octets_To_Char_Array( Key.n, N_char );
    Octets_To_Char_Array( Key.e, E_char );
    -- call C imported function
    out_len := Public_RSA_C( Encr_char  , Encr'Length, 
                             Plain_char , Plain'Length,
                             N_char     , Key.n'Length,
                             E_char     , Key.e'Length);
    -- convert back to octets
    Char_Array_To_Octets( Encr_char, Encr );
    -- C code trims leading 0s -> need to move octets if out_len<Encr'Length
    if out_len < Encr'Length then
      Encr(Encr'Last - out_len + 1 .. Encr'Last) := Encr(Encr'First .. 
                                                      Encr'First+out_len-1);
      Encr(Encr'First .. Encr'Last-out_len) := (others=>0);
    end if;
    -- no need to return anything!
  end Public_RSA;

  procedure Private_RSA( Encr  : in Raw_Types.RSA_len;
                         Key   : in RSA_skey;
                         Plain : out Raw_Types.Octets) is
    Plain_Char : char_array( size_t(Plain'First) .. size_t(Plain'Last) );
    Plain_Len  : Integer;
    Encr_Char  : char_array( size_t(Encr'First)  .. size_t(Encr'Last) );
    N_Char     : char_array( size_t(Key.n'First) .. size_t(Key.n'Last) );
    E_Char     : char_array( size_t(Key.e'First) .. size_t(Key.e'Last) );
    D_Char     : char_array( size_t(Key.d'First) .. size_t(Key.d'Last) );
    P_Char     : char_array( size_t(Key.p'First) .. size_t(Key.p'Last) );
    Q_Char     : char_array( size_t(Key.q'First) .. size_t(Key.q'Last) );
    U_Char     : char_array( size_t(Key.u'First) .. size_t(Key.u'Last) );
  begin
    -- convert key and encrypted message to C char_arrays
    Octets_To_Char_Array( Encr, Encr_Char );
    Octets_To_Char_Array( Key.n, N_Char );
    Octets_To_Char_Array( Key.e, E_Char );
    Octets_To_Char_Array( Key.d, D_Char );
    Octets_To_Char_Array( Key.p, P_Char );
    Octets_To_Char_Array( Key.q, Q_Char );
    Octets_To_Char_Array( Key.u, U_Char );
    -- call RSA decrypt via C_Wrappers
    Plain_Len := Private_RSA_C( Plain_Char, Plain'Length,
                                Encr_Char , Encr'Length,
                                N_Char    , Key.n'Length,
                                E_Char    , Key.e'Length,
                                D_Char    , Key.d'Length,
                                P_Char    , Key.p'Length,
                                Q_Char    , Key.q'Length,
                                U_Char    , Key.u'Length);
    -- convert result back to Octets
    Char_Array_To_Octets( Plain_Char, Plain );
    -- if length < OAEP_Block'Length,it's 0-led and got trimmed, so move it
    if Plain_Len < Plain'Length then
      Plain(Plain'Last-Plain_Len+1..Plain'Last):= Plain(Plain'First ..
                                                   Plain'First + Plain_Len -1);
      Plain(Plain'First .. Plain'Last-Plain_Len) := (others => 0);
    end if;
    -- no need to return anything!
  end Private_RSA;

  procedure Octets_To_Char_Array( O   : in Raw_Types.Octets;
                                  A   : out Interfaces.C.char_array) is
  begin
    -- check that lengths ARE THE SAME!
    if A'Length /= O'Length then
      raise Mismatched_Lengths_Error;
    end if;
    -- copy values over octet by octet
    for I in 0 .. O'Length-1 loop
      A( A'First + Interfaces.C.size_t( I )) := 
                     Interfaces.C.Char( Character'Val(O(O'First + I)));
    end loop;
  end Octets_To_Char_Array;

  procedure Char_Array_To_Octets( A   : in Interfaces.C.char_array;
                                  O   : out Raw_Types.Octets) is
  begin
    -- check that lengths ARE THE SAME!
    if A'Length /= O'Length then
      raise Mismatched_Lengths_Error;
    end if;
    -- copy values over octet by octet
    for I in 0..O'Length -1 loop
      O( O'First + I ) := Character'Pos( Character(
                             A( A'First + Interfaces.C.size_t( I )) ));
    end loop;
  end Char_Array_To_Octets;

end RSA_OAEP;