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

(* TODO key format
   - check what is the exact format in GNUNET
   - endianess issue? *)
module EddsaPublicKey = 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 = 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 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
  (* TODO key format
     endianess issue? *)
  (* 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 "signature verification failure: invalid signature"
    | true -> Ok ()

  let to_octets t = t

  let of_octets v =
    match String.length v = 64 with
    | false ->
        Fmt.error "EddsaSignature.of_octets failure: data is not 64 bytes."
    | true -> Ok 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 check_size t =
    match String.length t = 64 with
    | false -> Error "EddsaSignature: invalid string length"
    | true -> Ok ()

  let jsont =
    let of_b32 s =
      let* t = B32.decode s in
      let+ () = check_size t in
      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
end = struct
  type t = string

  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 denomination_hash = Hash.DenominationHash.t

(* WIP *)
module FDH_RSA = struct
  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

  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 =
      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
      assert (String.length okm = len);
      let r = Mirage_crypto_pk.Z_extra.of_octets_be okm in
      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_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
    Mirage_crypto_pk.Z_extra.to_octets_be data_r_e

  (*
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
  *)
end