(* TODO bin
   - no [Bin.of_string] ? *)
let bin_of_string bin s =
  let v = Bin.decode bin s (ref 0) in
  Ok v

module Bin_rsa = struct
  (* TODO
     - need to strip leading zeros?
     - endianess ok?
     - tests *)
  (* 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 *)
  (* RSA private key custom format is inspired by the public rsa key format
     used by secmod to save private key to file *)
  open Syntax

  module Internal = struct
    let rev_string len s = String.init len (fun i -> s.[len - 1 - i])

    (* we need reverse bytes because Z.of_bits reads bytes in little endian *)
    let z_of_bits_be src pos len =
      String.sub src pos len |> rev_string len |> Z.of_bits

    let z_to_bits_be z =
      let bits = Z.to_bits z in
      rev_string (String.length bits) bits

    let check = function false -> Error (`Msg "invalid data") | true -> Ok ()

    let z_array_to_octets (arr : Z.t array) =
      let nb = Array.length arr in
      (*Fmt.epr "z_array_to_octets, nb=%d@." nb;*)
      let bits_arr = Array.map z_to_bits_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 =
      (*Fmt.epr "z_array_of_octets, nb=%d@." nb;*)
      let s_len = String.length s in
      let* () = check (s_len > 2 * nb) in
      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* () =
        let len = (2 * nb) + Array.fold_left ( + ) 0 len_arr in
        check (s_len = len)
      in
      let z_arr =
        Array.init nb (fun i ->
            let len = len_arr.(i) in
            let z = z_of_bits_be s !pos len in
            pos := !pos + len;
            z)
      in
      Ok z_arr
  end

  open Internal

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

  let pub_of_octets s =
    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 in
          Ok pub
      | _ -> assert false
    in
    match pub_of_octets s with
    | Error (`Msg e) -> Fmt.failwith "rsa pub_of_octets failure: %s@." e
    | Ok v -> v

  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 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' in
          Ok priv
      | _ -> assert false
    in
    match priv_of_octets s with
    | Error (`Msg e) -> Fmt.failwith "rsa priv_of_octets failure: %s@." e
    | Ok v -> v
end