module Spec.Chacha20Poly1305.Test
#reset-options "--z3rlimit 100 --initial_fuel 0 --max_fuel 0 --initial_ifuel 0"

open FStar.Mul
open Lib.IntTypes
open Lib.RawIntTypes
open Lib.Sequence
open Lib.ByteSequence
open Spec.Chacha20Poly1305

(* Tests: RFC7539 *)

let test_key = List.Tot.map u8_from_UInt8 [
         0x80uy;0x81uy;0x82uy;0x83uy;0x84uy;0x85uy;0x86uy;0x87uy;
         0x88uy;0x89uy;0x8auy;0x8buy;0x8cuy;0x8duy;0x8euy;0x8fuy;
	 0x90uy;0x91uy;0x92uy;0x93uy;0x94uy;0x95uy;0x96uy;0x97uy;
         0x98uy;0x99uy;0x9auy;0x9buy;0x9cuy;0x9duy;0x9euy;0x9fuy]

let test_nonce = List.Tot.map u8_from_UInt8 [
         0x07uy;0x00uy;0x00uy;0x00uy;0x40uy;0x41uy;0x42uy;0x43uy;
         0x44uy;0x45uy;0x46uy;0x47uy]

let test_plaintext = List.Tot.map u8_from_UInt8 [
         0x4cuy; 0x61uy; 0x64uy; 0x69uy; 0x65uy; 0x73uy; 0x20uy; 0x61uy;
    	 0x6euy; 0x64uy; 0x20uy; 0x47uy; 0x65uy; 0x6euy; 0x74uy; 0x6cuy;
    	 0x65uy; 0x6duy; 0x65uy; 0x6euy; 0x20uy; 0x6fuy; 0x66uy; 0x20uy;
    	 0x74uy; 0x68uy; 0x65uy; 0x20uy; 0x63uy; 0x6cuy; 0x61uy; 0x73uy;
    	 0x73uy; 0x20uy; 0x6fuy; 0x66uy; 0x20uy; 0x27uy; 0x39uy; 0x39uy;
    	 0x3auy; 0x20uy; 0x49uy; 0x66uy; 0x20uy; 0x49uy; 0x20uy; 0x63uy;
   	 0x6fuy; 0x75uy; 0x6cuy; 0x64uy; 0x20uy; 0x6fuy; 0x66uy; 0x66uy;
    	 0x65uy; 0x72uy; 0x20uy; 0x79uy; 0x6fuy; 0x75uy; 0x20uy; 0x6fuy;
    	 0x6euy; 0x6cuy; 0x79uy; 0x20uy; 0x6fuy; 0x6euy; 0x65uy; 0x20uy;
    	 0x74uy; 0x69uy; 0x70uy; 0x20uy; 0x66uy; 0x6fuy; 0x72uy; 0x20uy;
    	 0x74uy; 0x68uy; 0x65uy; 0x20uy; 0x66uy; 0x75uy; 0x74uy; 0x75uy;
    	 0x72uy; 0x65uy; 0x2cuy; 0x20uy; 0x73uy; 0x75uy; 0x6euy; 0x73uy;
    	 0x63uy; 0x72uy; 0x65uy; 0x65uy; 0x6euy; 0x20uy; 0x77uy; 0x6fuy;
    	 0x75uy; 0x6cuy; 0x64uy; 0x20uy; 0x62uy; 0x65uy; 0x20uy; 0x69uy;
    	 0x74uy; 0x2euy]

let test_aad = List.Tot.map u8_from_UInt8 [
           0x50uy;0x51uy;0x52uy;0x53uy;0xc0uy;0xc1uy;0xc2uy;0xc3uy;
           0xc4uy;0xc5uy;0xc6uy;0xc7uy]

let test_cipher = List.Tot.map u8_from_UInt8 [
               0xd3uy;0x1auy;0x8duy;0x34uy;0x64uy;0x8euy;0x60uy;0xdbuy;
               0x7buy;0x86uy;0xafuy;0xbcuy;0x53uy;0xefuy;0x7euy;0xc2uy;
	       0xa4uy;0xaduy;0xeduy;0x51uy;0x29uy;0x6euy;0x08uy;0xfeuy;
               0xa9uy;0xe2uy;0xb5uy;0xa7uy;0x36uy;0xeeuy;0x62uy;0xd6uy;
	       0x3duy;0xbeuy;0xa4uy;0x5euy;0x8cuy;0xa9uy;0x67uy;0x12uy;
               0x82uy;0xfauy;0xfbuy;0x69uy;0xdauy;0x92uy;0x72uy;0x8buy;
	       0x1auy;0x71uy;0xdeuy;0x0auy;0x9euy;0x06uy;0x0buy;0x29uy;
               0x05uy;0xd6uy;0xa5uy;0xb6uy;0x7euy;0xcduy;0x3buy;0x36uy;
	       0x92uy;0xdduy;0xbduy;0x7fuy;0x2duy;0x77uy;0x8buy;0x8cuy;
               0x98uy;0x03uy;0xaeuy;0xe3uy;0x28uy;0x09uy;0x1buy;0x58uy;
	       0xfauy;0xb3uy;0x24uy;0xe4uy;0xfauy;0xd6uy;0x75uy;0x94uy;
               0x55uy;0x85uy;0x80uy;0x8buy;0x48uy;0x31uy;0xd7uy;0xbcuy;
	       0x3fuy;0xf4uy;0xdeuy;0xf0uy;0x8euy;0x4buy;0x7auy;0x9duy;
               0xe5uy;0x76uy;0xd2uy;0x65uy;0x86uy;0xceuy;0xc6uy;0x4buy;
	       0x61uy;0x16uy;]

let test_mac = List.Tot.map u8_from_UInt8 [0x1auy;0xe1uy;0x0buy;0x59uy;0x4fuy;0x09uy;0xe2uy;0x6auy;
            0x7euy;0x90uy;0x2euy;0xcbuy;0xd0uy;0x60uy;0x06uy;0x91uy]

let test() =
  assert_norm(List.Tot.length test_key = 32);
  assert_norm(List.Tot.length test_nonce = 12);
  assert_norm(List.Tot.length test_plaintext = 114);
  assert_norm(List.Tot.length test_aad = 12);
  assert_norm(List.Tot.length test_cipher = 114);
  assert_norm(List.Tot.length test_mac = 16);
  let test_key: key = of_list test_key in
  let test_nonce:nonce = of_list test_nonce in
  let test_plaintext: lbytes 114 = of_list test_plaintext in
  let test_aad: lbytes 12 = of_list test_aad in
  let test_cipher : lbytes 114 = of_list test_cipher in
  let test_mac : lbytes 16 = of_list test_mac in
  let enc =  aead_encrypt test_key test_nonce test_plaintext test_aad in
  let cipher = Seq.slice enc 0 114 in
  let mac = Seq.slice enc 114 130 in
  let dec =  aead_decrypt test_key test_nonce test_cipher test_mac test_aad in
  let result_c = for_all2 (fun a b -> uint_to_nat #U8 a = uint_to_nat #U8 b) cipher test_cipher in
  let result_m = for_all2 (fun a b -> uint_to_nat #U8 a = uint_to_nat #U8 b) mac test_mac in
  let result_p = Some? dec && for_all2 (fun a b -> uint_to_nat #U8 a = uint_to_nat #U8 b) (Some?.v dec) test_plaintext in
  IO.print_string   "Expected cipher:";
  List.iter (fun a -> IO.print_string (UInt8.to_string (u8_to_UInt8 a))) (to_list test_cipher);
  IO.print_string "\nComputed cipher:";
  List.iter (fun a -> IO.print_string (UInt8.to_string (u8_to_UInt8 a))) (to_list cipher);
  IO.print_string   "\nExpected mac:";
  List.iter (fun a -> IO.print_string (UInt8.to_string (u8_to_UInt8 a))) (to_list test_mac);
  IO.print_string "\nComputed mac:";
  List.iter (fun a -> IO.print_string (UInt8.to_string (u8_to_UInt8 a))) (to_list mac);
  IO.print_string   "\nExpected plaintext:";
  List.iter (fun a -> IO.print_string (UInt8.to_string (u8_to_UInt8 a))) (to_list test_plaintext);
  IO.print_string "\nComputed plaintext:";
  if Some? dec then
    List.iter (fun a -> IO.print_string (UInt8.to_string (u8_to_UInt8 a))) (to_list (Some?.v dec));
  if result_c && result_m && result_p then begin IO.print_string "\n\nChacha20 : Success!\n"; true end
  else begin IO.print_string "\n\nChacha20: Failure :(\n"; false end