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 to_b32 t = B32.encode (to_octets t)

  let of_b32 s =
    let* s = B32.decode s in
    let+ v = of_octets s in
    v

  let jsont = 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 = 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

module DenominationHash : sig
  type t

  val bin : t Bin.t
  val caqti : t Caqti_type.t
  val jsont : t Jsont.t
  val hash_of_rsa : RsaPublicKey.t -> t
  val of_octets : string -> t
  val to_octets : t -> string
  val of_b32 : B32.t -> (t, string) result
  val to_b32 : t -> B32.t
end = struct
  open Digestif

  type t = SHA512.t

  let hash_of_rsa pub =
    let s = RsaPublicKey.to_octets pub in
    SHA512.digest_string s

  let of_octets s =
    match SHA512.of_raw_string_opt s with
    | None -> Fmt.failwith "H64.of_octets failure"
    | Some t -> t

  let to_octets = SHA512.to_raw_string
  let of_b32 s = Result.map of_octets (B32.decode s)
  let to_b32 t = B32.encode (to_octets t)

  let bin =
    let open Bin in
    map (bytes 64) of_octets to_octets

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

  let jsont = Jsont.of_of_string ~kind:"DenominationHash" 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 = DenominationHash.t