(* WIP fdh-rsa
   full-domain-hash RSA
   based on libgnunetutil crypto_rsa.c *)

module Z_extra = Mirage_crypto_pk.Z_extra
open Crypto

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

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

  let kdf_mod_n ~n ~xts ~ikm ~ctx =
    let nbits = Z.numbits n in
    let len = ((nbits - 1) / 8) + 1 in
    assert (8 * len = nbits);
    let rec go ctr =
      (* cat ctx ctr_be *)
      let ctx =
        let ctx_len = String.length ctx in
        let b = Bytes.create (ctx_len + 2) in
        Bytes.blit_string ctx 0 b 0 ctx_len;
        Bytes.set_uint16_be b ctx_len ctr;
        Bytes.unsafe_to_string b
      in
      let okm = kdf ~xts ~ikm ~ctx ~len in
      assert (String.length okm = len);
      let r = Z_extra.of_octets_be okm in
      if Z.gt r n then go (succ ctr) else r
    in
    go 0
end

let gcd_validate r n =
  match Z.equal (Z.gcd r n) Z.one with
  | true -> ()
  | false -> Fmt.failwith "RSA key is malicious"

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

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.kdf_mod_n ~n:pub.n ~xts ~ikm:bks ~ctx in
  gcd_validate r pub.n; r

let blind_msg 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_sec bkey pub.e pub.n in
  let data_r_e = Z.rem (Z.mul data r_e) pub.n in
  Z_extra.to_octets_be data_r_e

let unblind_sig pub ~bks bsig =
  let data = Z_extra.of_octets_be bsig in
  let bkey = rsa_blinding_key_derive pub bks in
  let r_inv =
    try Z.invert bkey pub.n
    with Division_by_zero ->
      (* => gcd(r,n) <> 1, should be already checked for *)
      assert false
  in
  let data = Z.rem (Z.mul data r_inv) pub.n in
  Z_extra.to_octets_be data

(* decrypt <=> sign *)
let sign ~key bmsg = Mirage_crypto_pk.Rsa.decrypt ~crt_hardening:true ~key bmsg

let verify ~key s ~msg =
  let msg_fdh = rsa_full_domain_hash key msg in
  let s1 = Z_extra.to_octets_be msg_fdh in
  let s2 = Mirage_crypto_pk.Rsa.encrypt ~key s in
  match Eqaf.equal s1 s2 with
  | false -> Fmt.error "RSA signature verification failed"
  | true -> Ok ()