open Syntax

module Binary_format_rsa = struct
  (* RSA public key binary format
       https://www.gnupg.org/documentation/manuals/gcrypt/MPI-formats.html
       := { uint16_be: n size; uint16_be: e size; n; e}

       integer in big-endian format (MSB first)
       leading zeroes are stripped unless they are required to keep a value positive
       no 0-termination *)

  let z_array_to_octets (arr : Z.t array) =
    let nb = Array.length arr in
    let bits_arr = Array.map Mirage_crypto_pk.Z_extra.to_octets_be arr in
    let len_arr = Array.map String.length bits_arr in
    let len = (2 * nb) + Array.fold_left ( + ) 0 len_arr in
    let b = Bytes.make len '\x00' in
    let pos = ref 0 in
    Array.iter
      (fun len ->
        Bytes.set_uint16_be b !pos len;
        pos := !pos + 2)
      len_arr;
    Array.iteri
      (fun i bits ->
        let len = len_arr.(i) in
        Bytes.blit_string bits 0 b !pos len;
        pos := !pos + len)
      bits_arr;
    Bytes.unsafe_to_string b

  let z_array_of_octets ~nb s =
    let s_len = String.length s in
    if s_len <= 2 * nb then Error "rsa of_octets error"
    else
      let pos = ref 0 in
      let len_arr =
        Array.init nb (fun _i ->
            let len = String.get_uint16_be s !pos in
            pos := !pos + 2;
            len)
      in
      let len = (2 * nb) + Array.fold_left ( + ) 0 len_arr in
      if s_len <> len then Error "rsa of_octets error"
      else
        let z_arr =
          Array.init nb (fun i ->
              let len = len_arr.(i) in
              let s = String.sub s !pos len in
              let z = Mirage_crypto_pk.Z_extra.of_octets_be s in
              pos := !pos + len;
              z)
        in
        Ok z_arr

  let pub_to_octets ({ n; e } : Mirage_crypto_pk.Rsa.pub) =
    z_array_to_octets [| n; e |]

  let pub_of_octets s =
    let* arr = z_array_of_octets ~nb:2 s in
    match arr with
    | [| n; e |] ->
        let+ pub = Mirage_crypto_pk.Rsa.pub ~n ~e |> unwrap_err_msg in
        pub
    | _ -> assert false

  (* custom private key binary format <> than gcrypt *)
  let priv_to_octets ({ e; d; n; p; q; dp; dq; q' } : Mirage_crypto_pk.Rsa.priv)
      =
    z_array_to_octets [| e; d; n; p; q; dp; dq; q' |]

  let priv_of_octets s =
    let* arr = z_array_of_octets ~nb:8 s in
    match arr with
    | [| e; d; n; p; q; dp; dq; q' |] ->
        let+ priv =
          Mirage_crypto_pk.Rsa.priv ~e ~d ~n ~p ~q ~dp ~dq ~q' |> unwrap_err_msg
        in
        priv
    | _ -> assert false
end

module EddsaPublicKey = struct
  open Mirage_crypto_ec.Ed25519

  type t = pub

  let to_octets t = pub_to_octets t

  let of_octets t =
    pub_of_octets t |> function
    | Error e -> Fmt.error "%a" Mirage_crypto_ec.pp_error e
    | Ok v -> Ok v

  let bin =
    let of_octets_exn t = of_octets t |> Result.get_ok in
    Bin.map (Bin.bytes 32) of_octets_exn to_octets

  let of_b32 s =
    let* octets = B32.decode s in
    let* pub = of_octets octets in
    Ok pub

  let to_b32 t = B32.encode (to_octets t)
  let jsont = Jsont.of_of_string ~kind:"EddsaPublicKey" of_b32 ~enc:to_b32

  let caqti =
    Caqti_type.custom
      ~encode:(fun v -> Ok (to_octets v))
      ~decode:(fun v -> of_octets v)
      Caqti_type.octets
end

module EddsaPrivateKey = struct
  (* EdDSA and ECDHE public keys always point on Curve25519
     and represented  using the standard 256 bits Ed25519 compact format,
     converted to Crockford Base32. *)
  open Mirage_crypto_ec.Ed25519

  type t = priv

  let generate = generate
  let pub_of_priv = pub_of_priv
  let to_octets t = priv_to_octets t

  let of_octets t =
    priv_of_octets t |> function
    | Error err ->
        let err = Fmt.str "%a" Mirage_crypto_ec.pp_error err in
        Error err
    | Ok v -> Ok v

  let bin =
    let of_octets_exn t = of_octets t |> Result.get_ok in
    Bin.map (Bin.bytes 32) of_octets_exn to_octets

  let jsont =
    let of_b32 s =
      let* octets = B32.decode s in
      of_octets octets
    in
    let to_b32 t = B32.encode (to_octets t) in
    Jsont.of_of_string ~kind:"EddsaPrivateKey" of_b32 ~enc:to_b32
end

module EddsaSignature : sig
  type t

  val sign : key:EddsaPrivateKey.t -> string -> t

  (* Ok () on verification success *)
  val verify : key:EddsaPublicKey.t -> t -> msg:string -> (unit, string) result
  val to_octets : t -> string
  val of_octets : string -> (t, string) result
  val jsont : t Jsont.t
  val bin : t Bin.t
  val caqti : t Caqti_type.t
end = struct
  (* transmitted as 64-bytes base32
     binary-encoded objects with just the R and S values *)
  type t = string

  (* mirage_crypto:
     "The result is the concatenation of r and s, as specified in RFC 8032." *)
  let sign ~key s = Mirage_crypto_ec.Ed25519.sign ~key s

  let verify ~key s ~msg =
    let b = Mirage_crypto_ec.Ed25519.verify ~key s ~msg in
    match b with
    | false -> Error "EddsaSignature verification: invalid signature"
    | true -> Ok ()

  let to_octets t = t

  let check_size t =
    match String.length t = 64 with
    | false -> Error "EddsaSignature of_octets: data is not 64 bytes."
    | true -> Ok ()

  let of_octets v =
    let+ () = check_size v in
    v

  let bin =
    let of_octets_exn t = of_octets t |> Result.get_ok in
    Bin.map (Bin.bytes 64) of_octets_exn to_octets

  let jsont =
    let of_b32 s =
      let* t = B32.decode s in
      of_octets t
    in
    let to_b32 = B32.encode in
    Jsont.of_of_string ~kind:"EddsaSignature" of_b32 ~enc:to_b32

  let caqti =
    Caqti_type.custom
      ~encode:(fun v -> Ok (to_octets v))
      ~decode:(fun s -> of_octets s)
      Caqti_type.octets
end

module RsaPublicKey = struct
  open Mirage_crypto_pk

  type t = Rsa.pub

  let to_octets = Binary_format_rsa.pub_to_octets
  let of_octets = Binary_format_rsa.pub_of_octets

  let jsont =
    let of_b32 s =
      let* s = B32.decode s in
      let+ v = of_octets s in
      v
    in
    let to_b32 t = B32.encode (to_octets t) in
    Jsont.of_of_string ~kind:"RsaPublicKey" of_b32 ~enc:to_b32

  let caqti : t Caqti_type.t =
    Caqti_type.custom
      ~encode:(fun v -> Ok (to_octets v))
      ~decode:(fun v -> of_octets v)
      Caqti_type.octets
end

module RsaPrivateKey = struct
  open Mirage_crypto_pk.Rsa

  type t = priv

  let generate ~bits () =
    let priv = generate ~bits () in
    let pub = pub_of_priv priv in
    (priv, pub)

  let pub_of_priv = pub_of_priv
  let of_octets = Binary_format_rsa.priv_of_octets
  let to_octets = Binary_format_rsa.priv_to_octets

  let jsont =
    let of_b32 s =
      let* s = B32.decode s in
      let+ v = of_octets s in
      v
    in
    let to_b32 t = B32.encode (to_octets t) in
    Jsont.of_of_string ~kind:"RsaPrivateKey" of_b32 ~enc:to_b32
end

module RsaSignature : sig
  type t

  val jsont : t Jsont.t
  val sign : key:RsaPrivateKey.t -> string -> t
end = struct
  type t = string

  (* TODO crypto
     this is a placeholder signature algorithm *)
  let sign ~key s = Mirage_crypto_pk.Rsa.PKCS1.sig_encode ~key s

  let jsont =
    let of_b32 s = B32.decode s in
    let to_b32 t = B32.encode t in
    Jsont.of_of_string ~kind:"RsaSignature" of_b32 ~enc:to_b32
end

(* some type aliases, just for prettier .mli *)
type eddsa_priv = EddsaPrivateKey.t
type eddsa_pub = EddsaPublicKey.t
type eddsa_sig = EddsaSignature.t
type rsa_priv = RsaPrivateKey.t
type rsa_pub = RsaPublicKey.t
type rsa_sig = RsaSignature.t
type denom_hash = Hash.DenominationHash.t

(* WIP *)
module FDH_RSA = struct
  open Mirage_crypto_pk

  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 rsa_blind pub ~bks ~msg =
    let data = rsa_full_domain_hash pub msg in
    let bkey = rsa_blinding_key_derive pub bks in
    (* can we just use [powm] here instead? *)
    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

  (* -- WIP crypto -- *)

  (* TODO crypto
     not sure about signature scheme used by taler
     libgnunetutil crypto_rsa.c use "(flags raw)" => no padding *)
  (* decrypt <=> sign *)
  let rsa_sign_z priv r =
    let data = Z_extra.to_octets_be r in
    Rsa.decrypt ~crt_hardening:true ~key:priv data

  (* TODO crypto
     look into mirage-crypto for this
     use Eqaf for constant time string compare *)
  let rsa_verify_z pub r sig_ =
    let data = Z_extra.to_octets_be r in
    let sig_' = Rsa.encrypt ~key:pub data in
    match String.equal sig_ sig_' with
    | false -> Fmt.error "RSA signature verification failed"
    | true -> Ok ()

  let rsa_sign_fdh priv msg =
    let pub = Rsa.pub_of_priv priv in
    let r = rsa_full_domain_hash pub msg in
    rsa_sign_z priv r

  let rsa_unblind pub ~bks ~sig_ =
    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 ubsig = Z.rem (Z.mul sig_ r_inv) pub.n in
    ubsig

  let rsa_verify pub ~msg ~sig_ =
    let r = rsa_full_domain_hash pub msg in
    rsa_verify_z pub r sig_
end