module Kdf = struct
  module XTR = Hkdf.Make (Digestif.SHA512)
  module PRF = Hkdf.Make (Digestif.SHA256)

  let kdf ~xts ~ikm ~ctx ~len =
    let prk = XTR.extract ~salt:xts ikm in
    let okm = PRF.expand ~prk ~info:ctx len in
    okm
end

(* WIP
let kdf_mod_n ~n ~xts ~ikm ~ctx =
  let nbits = Z.numbits n in
  let rec go ctr =
    let len = ((nbits - 1) / 8) + 1 in
    let ctr_be =
      let b = Bytes.create 2 in
      Bytes.set_uint16_be b 0 ctr;
      Bytes.unsafe_to_string b
    in
    let ctx = String.cat ctx ctr_be in
    let okm = Kdf.kdf ~xts ~ikm ~ctx ~len in
    let r = Z.of_bits okm in
    (* we do that because n is odd?
       TODO
       gcry_mpi_clear_highbit .. *)
    if Z.gt r n then go (succ ctr) else r
  in
  go 0

let rsa_full_domain_hash pub msg =
  let xts = RsaPublicKey.to_octets pub in
  let ctx = "RSA-FDA FTpsW!" in
  let r = kdf_mod_n ~n:pub.n ~xts ~ikm:msg ~ctx in
  r

let rsa_sign_fdh priv msg =
  let pub = Mirage_crypto_pk.Rsa.pub_of_priv priv in
  let v = rsa_full_domain_hash pub msg in
  let v = Z.to_bits v in
  Mirage_crypto_pk.Rsa.PKCS1.sig_encode ~crt_hardening:true ~key:priv v

let rsa_blinding_key_derive (pub : RsaPublicKey.t) bks =
  let xts = "Blinding KDF extractor HMAC key" in
  let ctx = "Blinding KDF" in
  let r = kdf_mod_n ~n:pub.n ~xts ~ikm:bks ~ctx in
  r

let rsa_blind pub ~bks ~msg =
  let data = rsa_full_domain_hash pub msg in
  let bkey = rsa_blinding_key_derive pub bks in
  let r_e = Z.powm bkey pub.e pub.n in
  let data_r_e = Z.rem (Z.mul data r_e) pub.n in
  Z.to_bits data_r_e
  *)