with SMG_OAEP; use SMG_OAEP;
with SMG_Keccak; use SMG_Keccak;
with Ada.Exceptions; use Ada.Exceptions;
with Ada.Text_IO; use Ada.Text_IO;
with Ada.Strings.Fixed; use Ada.Strings.Fixed;
with Interfaces; use Interfaces;

procedure SMG_Keccak.Test is
  --types
  type Keccak_Perms is (None, Theta, Rho, Pi, Chi, Iota);
  type Test_Vector is array(Keccak_Perms) of State;
  type Test_Round is array(Round_Index) of Test_Vector;

  --helper methods
  procedure print_state(S: in State; Title: in String) is
    Hex: array(0..15) of Character := ("0123456789ABCDEF");
    Len: constant Natural := Z_Length / 4;
    HexString: String(1..Len);
    W: ZWord;
  begin
    Put_Line("---------" & Title & "---------");
    for Y in XYCoord loop
      for X in XYCoord loop
        W := S(X,Y);
        for Z in 0..Len-1 loop
          HexString(Natural(Len-Z)) := Hex(Natural(W mod 16));
          W := W / 16;
        end loop;
        Put(HexString & " ");
      end loop;
      Put_Line("");
    end loop;
  end;

  procedure print_bitstream(B: in Bitstream; Title: in String) is
    Hex       : array(0..15) of Character := ("0123456789ABCDEF");
    HexString : String(1..B'Length/4);
    C         : Natural;
    Pos       : Natural;
  begin
    for I in 1..B'Length/4 loop
      Pos := (I-1)*4 + B'First;
      C := Natural( B(Pos) ) * 8 +
           Natural( B(Pos + 1) ) * 4 +
           Natural( B(Pos + 2) ) * 2 +
           Natural( B(Pos + 3) );
			HexString(I) := Hex(C);
    end loop;
    Put_Line("---" & Title & "---");
    Put_Line(HexString);
  end print_bitstream;

  function read_state(File: in FILE_TYPE; Oct: Positive :=8) return State is
    S: State;
    Line1: String := "0000000000000000 " &
                     "0000000000000000 " &
                     "0000000000000000 " & 
                     "0000000000000000 " &
                     "0000000000000000";
    StartPos, EndPos: Positive;
    Len: Positive := Oct*2;
  begin
    for Y in XYCoord loop
      Line1 := Get_Line(File);
      StartPos := Line1'First;
      EndPos := StartPos + Len-1;

      for X in XYCoord loop
        S(X,Y) := ZWord'value("16#" & Line1(StartPos..EndPos) & "#");
        StartPos := EndPos + 2;	--one space to skip
        EndPos := StartPos + Len - 1;
      end loop;
    end loop;
    return S;
  end read_state;

  --reads a full test round from specified file (pre-defined format)
  function read_from_file (filename : in String;
                           T        : out Test_Round)
                           return Boolean is 
    file: FILE_TYPE;
    InputMarker: String := "lanes as 64-bit words:";
    octets: Positive := 8;
    RoundNo: Round_Index;
  begin
    -- try to open the input file
    begin
      open(file, In_File, filename);
    exception
      when others =>
        Put_Line(Standard_Error, 
                 "Can not open the file '" & filename & "'. Does it exist?");
        return False;
    end;

  -- find & read input state first
    RoundNo := -1;
    loop
      declare
        Line: String := Get_Line(file);
      begin
        --check if this is test data of any known kind
        if index(Line, InputMarker, 1) > 0 then
          T(0)(None) := read_state(file, octets);
          print_state(T(0)(None), "Read Input State");
        elsif index(Line, "Round ", 1) > 0 then
          RoundNo := RoundNo +1;
        elsif index(Line, "theta", 1) > 0 then
          T(RoundNo)(Theta) := read_state(file, octets);
          if (RoundNo > 0) then
            T(RoundNo)(None) := T(RoundNo-1)(Iota);  -- previous state as input
          end if;
        elsif index(Line, "rho", 1) > 0 then
          T(RoundNo)(Rho) := read_state(file, octets);
        elsif index(Line, "pi", 1) > 0 then
          T(RoundNo)(Pi) := read_state(file, octets);
        elsif index(Line, "chi", 1) > 0 then
          T(RoundNo)(Chi) := read_state(file, octets);
        elsif index(Line, "iota", 1) > 0 then
          T(RoundNo)(Iota) := read_state(file, octets);
        end if;	
        exit when End_Of_File(file);
      end;
    end loop;
    Close(file);
    return True;
  end read_from_file;

  -- performs one single round of Keccak, step by step
  -- each permutation is tested separately
  -- test fails with exception raised at first output not matching expected
  procedure test_one_round(T: Test_Vector; Round: Round_Index) is
    Input: State;
    Expected: State;
    Output: State;
    Test_One_Round_Fail: Exception;
  begin
    Input := T(None);
    for I in Keccak_Perms range Theta..Iota loop
      Expected := T(I);
      case I is
        when Theta => Output := SMG_Keccak.Theta(Input);
        when Rho   => Output := SMG_Keccak.Rho(Input);
        when Pi    => Output := SMG_Keccak.Pi(Input);
        when Chi => Output := SMG_Keccak.Chi(Input);
        when Iota => Output := SMG_Keccak.Iota(RC(Round), Input);
        when others => null;
      end case;

      if (Output /= Expected) then
        print_state(Output, "----------real output-------");
        print_state(Expected, "----------expected output--------");
        raise Test_One_Round_Fail;
      else
        Put_Line("PASSED: " & Keccak_Perms'Image(I));
      end if;
      -- get ready for next permutation
      Input := Expected;
    end loop;
  end test_one_round;

  procedure test_bits_to_word_conversion is
    bits: Bitword := (others => 0);
    obtained_bits: Bitword := (others => 0);
    expected: ZWord;
    obtained: ZWord;
  begin
    expected := 16#8FA4F19E0287BBE7#;
    bits := (1,1,1,0, 0,1,1,1, 1,1,0,1, 1,1,0,1, 1,1,1,0, 0,0,0,1, 0,1,0,0, 
             0,0,0,0, 0,1,1,1, 1,0,0,1, 1,0,0,0, 1,1,1,1, 0,0,1,0, 0,1,0,1, 
             1,1,1,1, 0,0,0,1);
    obtained := BitsToWord(bits);
    obtained_bits := WordToBits(expected);

    if obtained /= expected then
      Put_Line("FAIL: bits to word");
      Put_Line("Expected: " & ZWord'Image(expected));
      Put_Line("Obtained: " & ZWord'Image(obtained));
    else
      Put_Line("PASSED: bits to word");
    end if;

    if obtained_bits /= bits then
      Put_Line("FAIL: word to bits");
      Put("Expected: ");
      for I in Bitword'Range loop
        Put(Bit'Image(bits(I)));
      end loop;
      Put_Line("");
      Put_Line("Obtained: ");
      for I in Bitword'Range loop
        Put(Bit'Image(obtained_bits(I)));
      end loop;
      Put_Line("");
    else
      Put_Line("PASSED: word to bits");
    end if;
  end test_bits_to_word_conversion;

  procedure test_flip is
    B: constant Bitword := (1, 0, 1, 1, 1, 1, 0, 0, 
                            1, 1, 1, 0, 0, 0, 0, 1,
                            0, 1, 1, 0, 0, 0, 1, 0,
                            1, 1, 1, 1, 1, 1, 1, 1,
                            1, 1, 0, 1, 1, 0, 0, 1, 
                            0, 0, 0, 0, 0, 0, 0, 0,
                            0, 0, 1, 1, 0, 0, 0, 1,
                            0, 0, 0, 1, 1, 0, 0, 0);
    Expected: Bitword :=   (0, 0, 0, 1, 1, 0, 0, 0,
                            0, 0, 1, 1, 0, 0, 0, 1,
                            0, 0, 0, 0, 0, 0, 0, 0,
                            1, 1, 0, 1, 1, 0, 0, 1, 
                            1, 1, 1, 1, 1, 1, 1, 1,
                            0, 1, 1, 0, 0, 0, 1, 0,
                            1, 1, 1, 0, 0, 0, 0, 1,
                            1, 0, 1, 1, 1, 1, 0, 0);
    Output : Bitword;
  begin
    Output := FlipOctets( B );
    if Output /= Expected then
      Put_Line( "FAILED: flip octets" );
      Put_Line( "Expected: " );
      for I of Expected loop
        Put(Bit'Image(I));
      end loop;
      new_line(1);
      Put_Line( "Output: " );
      for I of Output loop
        Put(Bit'Image(I));
      end loop;
      new_line(1);
    else
      Put_Line( "PASSED: flip octets" );
    end if;
  end test_flip;

  procedure test_sponge is
    Bitrate   : constant Keccak_Rate := 1344;
    Input     : Bitstream(1..5) := (1, 1, 0, 0, 1);
    Hex       : array(0..15) of Character := ("0123456789ABCDEF");
    C         : Natural;
    HexPos    : Natural;
    Error     : Natural;
    Pos       : Natural;
    ExpHex    : constant String := 
              "CB7FFB7CE7572A06C537858A0090FC2888C3C6BA9A3ADAB4"&
              "FE7C9AB4EFE7A1E619B834C843A5A79E23F3F7E314AA597D"&
              "9DAD376E8413A005984D00CF954F62F59EF30B050C99EA64"&
              "E958335DAE684195D439B6E6DFD0E402518B5E7A227C48CF"&
              "239CEA1C391241D7605733A9F4B8F3FFBE74EE45A40730ED"&
              "1E2FDEFCCA941F518708CBB5B6D5A69C30263267B97D7B29"&
              "AC87043880AE43033B1017EFB75C33248E2962892CE69DA8"&
              "BAF1DF4C0902B16C64A1ADD42FF458C94C4D3B0B32711BBA"&
              "22104989982543D1EF1661AFAF2573687D588C81113ED7FA"&
              "F7DDF912021FC03D0E98ACC0200A9F7A0E9629DBA33BA0A3"&
              "C03CCA5A7D3560A6DB589422AC64882EF14A62AD9807B353"&
              "8DEE1548194DBD456F92B568CE76827F41E0FB3C7F25F3A4"&
              "C707AD825B289730FEBDFD22A3E742C6FB7125DE0E38B130"&
              "F3059450CA6185156A7EEE2AB7C8E4709956DC6D5E9F99D5"&
              "0A19473EA7D737AC934815D68C0710235483DB8551FD8756"&
              "45692B4E5E16BB9B1142AE300F5F69F43F0091D534F372E1"&
              "FFC2E522E71003E4D27EF6ACCD36B2756FB5FF02DBF0C96B"&
              "CAE68E7D6427810582F87051590F6FB65D7B948A9C9D6C93"&
              "AF4562367A0AD79109D6F3087C775FE6D60D66B74F8D29FB"&
              "4BA80D0168693A748812EA0CD3CA23854CC84D4E716F4C1A"&
              "A3B340B1DED2F304DFDBACC1D792C8AC9A1426913E3F67DB"& 
              "790FD5CFB77DAA29"; 
    Output    : Bitstream( 1 .. ExpHex'Length * 4 );
    HexString : String( 1 .. ExpHex'Length );
  begin
    Put_Line("---sponge test---");
    Sponge(Input, Output, Bitrate);
    Put_Line("Input is:");
    for I of Input loop
      Put(Bit'Image(I));
    end loop;
    new_line(1);

    Put_Line("Output is:");
    for I of Output loop
      Put(Bit'Image(I));
    end loop;
    new_line(1);

    Error := 0;
    for I in 1..Output'Length/4 loop
      Pos := Output'First + (I-1)*4;
      C := Natural( Output( Pos ) ) +
           Natural( Output( Pos + 1 ) ) * 2 +
           Natural( Output( Pos + 2 ) ) * 4 +
           Natural( Output( Pos + 3 ) ) * 8;
      HexPos := I + 2 * ( I mod 2 ) - 1; 
      Hexstring(HexPos) := Hex( C );
      if Hexstring(HexPos) /= ExpHex(HexPos) then
        Error := Error + 1;
      end if;
    end loop;
    Put_Line("Expected: ");
    Put_Line(ExpHex);
    Put_Line("Obtained: ");
    Put_Line(Hexstring);
    Put_Line("Errors found: " & Natural'Image(Error));

  end test_sponge;

  procedure test_keccak_function(T: in Test_Round) is
    S: State;
  begin
    Put_Line("---Full Keccak Function test---");
    S := Keccak_Function(T(Round_Index'First)(None));
    if S /= T(Round_Index'Last)(Iota) then
      Put_Line("FAILED: full keccak function test");
    else
      Put_Line("PASSED: full keccak function test");
    end if;
  end test_keccak_function;

  procedure test_bitstream_conversion is
    S: String := "Aa*/";
    E: Bitstream( 0 .. 31 ) := (1, 0, 0, 0, 0, 0, 1, 0,
                                1, 0, 0, 0, 0, 1, 1, 0,
                                0, 1, 0, 1, 0, 1, 0, 0,
                                1, 1, 1, 1, 0, 1, 0, 0);
    B: Bitstream( 0 .. 31 );
    SS: String := "  t ";
  begin
    Put_Line("---Testing string to bitstream conversion---");
    ToBitstream( S, B );
    if E /= B then
      Put_Line("FAILED: string to bitstream conversion.");
    else
      Put_Line("PASSED: string to bitstream conversion.");
    end if;

    Put_Line("---Testing bitstream to string conversion---");
    ToString( B, SS );
    if SS /= S then
      Put_Line("FAILED: bitstream to string conversion");
      Put_Line("EXPECTED: " & S);
      Put_Line("OUTPUT: " & SS);
    else
      Put_Line("PASSED: bitstream to string conversion");
    end if;
  end test_bitstream_conversion;

  procedure test_hash_keccak is
    S: String := "X";
    O: String := "abc";
    B: Bitstream( 0 .. 23 );
    BB: Bitstream( 1.. 8):= (0, 0, 0, 1, 1, 0, 1, 0);
    Exp: Bitstream( 0 .. 23 ) := (1, 1, 1, 0, 0, 0, 0, 1,
                                  0, 1, 1, 0, 0, 0, 1, 0,
                                  1, 1, 1, 0, 0, 0, 1, 1);
  begin
    Put_Line("----Testing hash keccak on string " & S & "----");
    HashKeccak(S, O);
    ToBitstream( O, B );
    if B /= Exp then
      Put_Line("FAILED: testing hash keccak on string");
      Put_Line("Output:");
      for I of B loop
        Put( Bit'Image( I ) );
      end loop;
      new_line(1);
      Put_Line("Expected:");
      for I of Exp loop
        Put( Bit'Image( I ) );
      end loop;
    else
      Put_Line("PASSED: testing hash keccak on string");
    end if;
    new_line(1);
  end test_hash_keccak;

  procedure test_xor_strings is
    S1     : String := "ABC";
    S2     : String := "CBA";
    Exp    : String := "...";
    Result : String := "...";
  begin
    Exp( Exp'First     ) := Character'Val( 2 );
    Exp( Exp'First + 1 ) := Character'Val( 0 );
    Exp( Exp'First + 2 ) := Character'Val( 2 );

    Put_Line("----Testing xor on strings---");
    XOR_Strings( S1, S2, Result);
    Put_Line("S1 is " & S1);
    Put_Line("S2 is " & S2);
    Put_Line("Result is: ");
    for C of Result loop
      Put( Natural'Image( Character'Pos( C ) ) );
    end loop;
    new_line(1);

    if Result /= Exp then
      Put_Line("FAILED: xor on strings");
    else
      Put_Line("PASSED: xor on strings");
    end if;
  end test_xor_strings;

  procedure test_oaep is
    Msg     : String := "abcdefghij jihgfedcba123456789";
    LongMsg : String( 1..1000 ) := ( others => 'T' );
    Encr    : OAEP_Block := ( others => ' ' );
    Decr    : OAEP_HALF  := ( others => ' ' );
    Entropy : OAEP_Block := ( others => 'e' );
    Len     : Natural;
    Flag    : Boolean;
    C       : Character;
    MaxLen  : constant := 245;
  begin
    Put_Line("----Testing OAEP Encrypt----");
    OAEP_Encrypt( Msg, Entropy, Encr );

    Put_Line("----Testing OAEP Decrypt----");
    OAEP_Decrypt( Encr, Len, Decr, Flag );

    if Flag = False or 
       Len /= Msg'Length * 8 or 
       Decr( Decr'First .. Decr'First + Msg'Length - 1 ) /= Msg
       then
      Put_Line("FAILED: oaep test");
      Put_Line("Msg is: "  & Msg);
      Put_Line("Decr is: " & Decr);
      Put_Line("Flag is: " & Boolean'Image( Flag ) );
      Put_Line("Len is: "  & Natural'Image( Len ) );
    else
      Put_Line("PASSED: oaep test");
    end if;

    -- test decrypt on invalid (non-OAEP) string
    Flag := True;
    C := Encr( Encr'First );
    Encr( Encr'First ) := Character'Val( Character'Pos( C ) / 2 );
    Decr := ( others => ' ' );
    OAEP_Decrypt( Encr, Len, Decr, Flag );

    if Flag = True then
      Put_Line("FAILED: oaep test with invalid package");
    else
      Put_Line("PASSED: oaep test with invalid package");
    end if; 

    -- test encrypt on message longer than maximum payload (1096 bits)
    Flag := False;
    Len := 0;
    LongMsg( 1..Msg'Length ) := Msg;
    Encr := ( others => '.' );
    OAEP_Encrypt( LongMsg, Entropy, Encr);
    OAEP_Decrypt( Encr, Len, Decr, Flag);

    if Flag = False or 
       Len /= MaxLen * 8 or
       Decr( Decr'First .. Decr'First + Len / 8 - 1 ) /= 
             LongMsg( LongMsg'First..LongMsg'First + MaxLen - 1 ) 
       then
      Put_Line("FAILED: oaep test with too long message");
      Put_Line("Msg is: "  & LongMsg);
      Put_Line("Decr is: " & Decr);
      Put_Line("Flag is: " & Boolean'Image( Flag ) );
      Put_Line("Len is: "  & Natural'Image( Len ) );
    else
      Put_Line("PASSED: oaep test with too long message");
    end if;

  end test_oaep;

  procedure test_all_bitrates is
    Input : constant String := "hello, world";
    Bin   : Bitstream( 0 .. Input'Length * 8 - 1 ) := ( others => 0 );
    Bout  : Bitstream( 0 .. 100 ) := ( others => 0 );
  begin
    ToBitstream( Input, Bin );
    Put_Line("Testing all bitrates:");
      for Bitrate in Keccak_Rate'Range loop
        Sponge(Bin, Bout, Bitrate);
        Put_Line("PASSED: keccak with bitrate " & Integer'Image(Bitrate));
      end loop;
  end test_all_bitrates;

  -- end of helper methods 

	--variables
  T: Test_Round;
begin
  Put_Line("-----Testing with zero state as input------");
  if (not read_from_file("testvectorszero.txt", T)) then
    return;
  end if;

  for I in Round_Index loop
    Put_Line("---round " & Round_Index'Image(I) & "---");
    test_one_round(T(I), I);		
  end loop;

  -- test also Keccak_Function as a whole --
  test_keccak_function(T);

  Put_Line("-----Testing with non-zero state as input------");
  if (not read_from_file("testvectorsnonzero.txt", T)) then 
    return;
  end if;

  for I in Round_Index loop
    Put_Line("---round " & Round_Index'Image(I) & "---");
    test_one_round(T(I), I);		
  end loop;

  -- test also Keccak_Function as a whole --
  test_keccak_function(T);

  -- test BitsToWord and WordToBits
  test_bits_to_word_conversion;

  -- test Sponge construction
  test_sponge;

  -- test flipping octets
  test_flip;

  -- test bitstream conversion
  test_bitstream_conversion;

  -- test hash keccak (strings version)
  test_hash_keccak;

  -- test oaep encrypt + decrypt
  test_oaep;

  -- test xor on strings
  test_xor_strings;

  -- test ALL bitrates of the Keccak sponge
  test_all_bitrates;

end SMG_Keccak.Test;