open Swim.Crypto

let valid_key = String.make 16 '\x00'

let test_roundtrip_property random =
  QCheck.Test.make ~count:500 ~name:"crypto roundtrip" Generators.arb_cstruct
    (fun plaintext ->
      match init_key valid_key with
      | Error _ -> false
      | Ok key -> (
          let ciphertext = encrypt ~key ~random plaintext in
          match decrypt ~key ciphertext with
          | Ok decrypted -> Cstruct.equal plaintext decrypted
          | Error _ -> false))

let test_encrypt_increases_size random =
  QCheck.Test.make ~count:100 ~name:"encrypt increases size by overhead"
    Generators.arb_cstruct (fun plaintext ->
      match init_key valid_key with
      | Error _ -> false
      | Ok key ->
          let ciphertext = encrypt ~key ~random plaintext in
          Cstruct.length ciphertext = Cstruct.length plaintext + overhead)

let test_different_plaintexts_different_ciphertexts random =
  QCheck.Test.make ~count:100
    ~name:"different plaintexts produce different ciphertexts"
    (QCheck.pair Generators.arb_cstruct Generators.arb_cstruct) (fun (p1, p2) ->
      if Cstruct.equal p1 p2 then true
      else
        match init_key valid_key with
        | Error _ -> false
        | Ok key ->
            let c1 = encrypt ~key ~random p1 in
            let c2 = encrypt ~key ~random p2 in
            not (Cstruct.equal c1 c2))

let test_init_key_valid_length () =
  match init_key (String.make 16 'a') with
  | Ok _ -> ()
  | Error _ -> Alcotest.fail "expected valid key"

let test_init_key_15_bytes_rejected () =
  match init_key (String.make 15 'a') with
  | Error `Invalid_key_length -> ()
  | _ -> Alcotest.fail "expected Invalid_key_length"

let test_init_key_17_bytes_rejected () =
  match init_key (String.make 17 'a') with
  | Error `Invalid_key_length -> ()
  | _ -> Alcotest.fail "expected Invalid_key_length"

let test_init_key_empty_rejected () =
  match init_key "" with
  | Error `Invalid_key_length -> ()
  | _ -> Alcotest.fail "expected Invalid_key_length"

let test_tampered_ciphertext_fails random () =
  match init_key valid_key with
  | Error _ -> Alcotest.fail "key init failed"
  | Ok key -> (
      let plaintext = Cstruct.of_string "hello world" in
      let ciphertext = encrypt ~key ~random plaintext in
      let tampered = Cstruct.of_string (Cstruct.to_string ciphertext) in
      let pos = Cstruct.length tampered - 1 in
      Cstruct.set_uint8 tampered pos
        ((Cstruct.get_uint8 tampered pos + 1) land 0xFF);
      match decrypt ~key tampered with
      | Error `Decryption_failed -> ()
      | _ -> Alcotest.fail "expected Decryption_failed")

let test_truncated_ciphertext_fails random () =
  match init_key valid_key with
  | Error _ -> Alcotest.fail "key init failed"
  | Ok key -> (
      let plaintext = Cstruct.of_string "hello world" in
      let ciphertext = encrypt ~key ~random plaintext in
      let truncated = Cstruct.sub ciphertext 0 (overhead - 1) in
      match decrypt ~key truncated with
      | Error `Too_short -> ()
      | _ -> Alcotest.fail "expected Too_short")

let test_wrong_key_fails random () =
  match (init_key valid_key, init_key (String.make 16 '\xFF')) with
  | Ok key1, Ok key2 -> (
      let plaintext = Cstruct.of_string "secret message" in
      let ciphertext = encrypt ~key:key1 ~random plaintext in
      match decrypt ~key:key2 ciphertext with
      | Error `Decryption_failed -> ()
      | _ -> Alcotest.fail "expected Decryption_failed")
  | _ -> Alcotest.fail "key init failed"

let test_empty_plaintext random () =
  match init_key valid_key with
  | Error _ -> Alcotest.fail "key init failed"
  | Ok key -> (
      let plaintext = Cstruct.empty in
      let ciphertext = encrypt ~key ~random plaintext in
      Alcotest.(check int)
        "ciphertext size" overhead
        (Cstruct.length ciphertext);
      match decrypt ~key ciphertext with
      | Ok decrypted ->
          Alcotest.(check int) "decrypted size" 0 (Cstruct.length decrypted)
      | Error _ -> Alcotest.fail "decrypt failed")

let test_nonce_uniqueness random () =
  match init_key valid_key with
  | Error _ -> Alcotest.fail "key init failed"
  | Ok key ->
      let plaintext = Cstruct.of_string "test" in
      let c1 = encrypt ~key ~random plaintext in
      let c2 = encrypt ~key ~random plaintext in
      let nonce1 = Cstruct.sub c1 1 nonce_size in
      let nonce2 = Cstruct.sub c2 1 nonce_size in
      if Cstruct.equal nonce1 nonce2 then
        Alcotest.fail "nonces should be different"
      else ()

let test_ciphertext_differs_from_plaintext random () =
  match init_key valid_key with
  | Error _ -> Alcotest.fail "key init failed"
  | Ok key ->
      let plaintext = Cstruct.of_string "hello world secret" in
      let ciphertext = encrypt ~key ~random plaintext in
      let ciphertext_body =
        Cstruct.sub ciphertext (1 + nonce_size)
          (Cstruct.length ciphertext - 1 - nonce_size)
      in
      if
        Cstruct.equal plaintext
          (Cstruct.sub ciphertext_body 0
             (min (Cstruct.length plaintext) (Cstruct.length ciphertext_body)))
      then Alcotest.fail "ciphertext should differ from plaintext"
      else ()

let () =
  Eio_main.run @@ fun env ->
  let random = Eio.Stdenv.secure_random env in
  let qcheck_tests =
    List.map QCheck_alcotest.to_alcotest
      [
        test_roundtrip_property random;
        test_encrypt_increases_size random;
        test_different_plaintexts_different_ciphertexts random;
      ]
  in
  let unit_tests =
    [
      ("init_key_valid_length", `Quick, test_init_key_valid_length);
      ("init_key_15_bytes_rejected", `Quick, test_init_key_15_bytes_rejected);
      ("init_key_17_bytes_rejected", `Quick, test_init_key_17_bytes_rejected);
      ("init_key_empty_rejected", `Quick, test_init_key_empty_rejected);
      ( "tampered_ciphertext_fails",
        `Quick,
        test_tampered_ciphertext_fails random );
      ( "truncated_ciphertext_fails",
        `Quick,
        test_truncated_ciphertext_fails random );
      ("wrong_key_fails", `Quick, test_wrong_key_fails random);
      ("empty_plaintext", `Quick, test_empty_plaintext random);
      ("nonce_uniqueness", `Quick, test_nonce_uniqueness random);
      ( "ciphertext_differs_from_plaintext",
        `Quick,
        test_ciphertext_differs_from_plaintext random );
    ]
  in
  Alcotest.run "crypto" [ ("property", qcheck_tests); ("unit", unit_tests) ]