-- Tests for SMG_Comms.Packing
-- S.MG, 2018

with Packing; use Packing;
with Raw_Types; use Raw_Types;
with RSA_OAEP; use RSA_OAEP;
with Serpent; use Serpent;
with RNG; use RNG; 
with IO_RSA;

with Interfaces; use Interfaces;
with Ada.Text_IO; use Ada.Text_IO;

package body Test_Packing is

  procedure Print(Data: in Raw_Types.Octets; Msg: in String) is
  begin
    Put_Line(Msg);
    for I of Data loop
      Put(Unsigned_8'Image(I));
    end loop;
    New_Line;
  end Print;

  procedure Test_Pack_Unpack_Serpent is
    InMsg  : Serpent_Msg := (others => 0);
    OutMsg : Serpent_Msg := (others => 0);

    InPkt  : Serpent_Pkt := (others => 0);
    OutPkt : Serpent_Pkt := (others => 0);

    K      : Key := (others => 0);
    KS     : Key_Schedule;
    Plain  : Block := (others => 0);
    Encr   : Block := (others => 0);
    Len    : constant Natural := 
             Raw_Types.SERPENT_OCTETS / Serpent.Block'Length;
  begin
    for I in 1 .. 128 loop
      OutPkt := Pack(InMsg, K);
      if OutPkt = InMsg then
        raise Test_Error;
      end if;

      OutMsg := Unpack(OutPkt, K);
      if OutMsg /= InMsg then
        raise Test_Error;
      end if;

      -- check result of pack
      Prepare_Key(K, KS);
      for J in 1 .. Len loop
        Plain := InMsg((J-1)*Block'Length + 1 .. J*Block'Length);
        Serpent.Encrypt(KS, Plain, Encr);
        if Encr /= OutPkt((J-1)*Block'Length + 1 .. J*Block'Length) then
          Put_Line("FAIL: pack/unpack with Serpent.");
          raise Test_Error;
        end if;
      end loop;
 
      -- iterate, re-packing as "message" the previous package
      InMsg := OutPkt;
    end loop;

    Put_Line("PASS: test pack/unpack with Serpent.");
  end Test_Pack_Unpack_Serpent;

  procedure Test_Pack_Unpack_RSA is
    Msg       : RSA_Msg;
    Decr_Msg  : RSA_Msg;
    PKey      : RSA_pkey;
    SKey      : RSA_skey;
    Success   : Boolean;
    Pkt       : RSA_Pkt;
  begin
    -- initialize with RSA pair previously generated
    IO_RSA.ReadRSAKey( "keys_rsa.txt", SKey );

    -- copy n and e for public key
    PKey.n := SKey.n;
    PKey.e := SKey.e;
    -- get random data for "message"
    RNG.Get_Octets(Msg);

    -- pack the message into corresponding packet
    Pkt := Packing.Pack( Msg, PKey );
    -- unpack and check the result
    Decr_Msg := Packing.Unpack( Pkt, SKey, Success );
    if (not Success) or (Decr_Msg /= Msg) then
      Put_Line("FAIL: pack/unpack with RSA.");
    else
      Put_Line("PASS: pack/unpack with RSA.");
    end if;

    -- try to unpack a mangled package
    Pkt(Pkt'First + 3) := Pkt(Pkt'First + 3) + 1;
    Decr_Msg := (others => 0);
    Decr_Msg := Packing.Unpack( Pkt, SKey, Success );
    if Success then
      Put_Line("FAIL: pack/unpack with RSA on mangled package.");
    else
      Put_Line("PASS: pack/unpack with RSA on mangled package.");
    end if;

  end Test_Pack_Unpack_RSA;

end Test_Packing;