open Uncommon

let block_size = 16

let flags bit6 len1 len2 =
  bit6 lsl 6 + len1 lsl 3 + len2

let encode_len buf ~off size value =
  let rec ass num = function
    | 0 -> Bytes.set_uint8 buf off num
    | m ->
      Bytes.set_uint8 buf (off + m) (num land 0xff);
      (ass [@tailcall]) (num lsr 8) (pred m)
  in
  ass value (pred size)

let set_format buf ?(off = 0) nonce flag_val value =
  let n = String.length nonce in
  let small_q = 15 - n in
  (* first octet block:
     0          : flags
     1..15 - q  : N
     16 - q..15 : Q *)
  Bytes.set_uint8 buf off flag_val;
  Bytes.unsafe_blit_string nonce 0 buf (off + 1) n;
  encode_len buf ~off:(off + n + 1) small_q value

let gen_adata a =
  let llen, set_llen =
    match String.length a with
    | x when x < (1 lsl 16 - 1 lsl 8) ->
      2, (fun buf off -> Bytes.set_uint16_be buf off x)
    | x when Sys.int_size < 32 || x < (1 lsl 32) ->
      6, (fun buf off ->
          Bytes.set_uint16_be buf off 0xfffe;
          Bytes.set_int32_be buf (off + 2) (Int32.of_int x))
    | x ->
      10, (fun buf off ->
          Bytes.set_uint16_be buf off 0xffff;
          Bytes.set_int64_be buf (off + 2) (Int64.of_int x))
  in
  let to_pad =
    let leftover = (llen + String.length a) mod block_size in
    block_size - leftover
  in
  llen + String.length a + to_pad,
  fun buf off ->
    set_llen buf off;
    Bytes.unsafe_blit_string a 0 buf (off + llen) (String.length a);
    Bytes.unsafe_fill buf (off + llen + String.length a) to_pad '\000'

let gen_ctr nonce i =
  let n = String.length nonce in
  let small_q = 15 - n in
  let flag_val = flags 0 0 (small_q - 1) in
  let buf = Bytes.create 16 in
  set_format buf nonce flag_val i;
  buf

let prepare_header nonce adata plen tlen =
  let small_q = 15 - String.length nonce in
  let b6 = if String.length adata = 0 then 0 else 1 in
  let flag_val = flags b6 ((tlen - 2) / 2) (small_q - 1) in
  if String.length adata = 0 then
    let hdr = Bytes.create 16 in
    set_format hdr nonce flag_val plen;
    hdr
  else
    let len, set = gen_adata adata in
    let buf = Bytes.create (16 + len) in
    set_format buf nonce flag_val plen;
    set buf 16;
    buf

type mode = Encrypt | Decrypt

let crypto_core_into ~cipher ~mode ~key ~nonce ~adata src ~src_off dst ~dst_off len =
  let cbcheader = prepare_header nonce adata len block_size in

  let small_q = 15 - String.length nonce in
  let ctr_flag_val = flags 0 0 (small_q - 1) in
  let ctrblock i block dst_off =
    Bytes.set_uint8 block dst_off ctr_flag_val;
    Bytes.unsafe_blit_string nonce 0 block (dst_off + 1) (String.length nonce);
    encode_len block ~off:(dst_off + String.length nonce + 1) small_q i;
    cipher ~key (Bytes.unsafe_to_string block) ~src_off:dst_off block ~dst_off
  in

  let cbc iv src_off block dst_off =
    unsafe_xor_into iv ~src_off block ~dst_off block_size ;
    cipher ~key (Bytes.unsafe_to_string block) ~src_off:dst_off block ~dst_off
  in

  let iv =
    let rec doit iv iv_off block block_off =
      match Bytes.length block - block_off with
      | 0 -> Bytes.sub iv iv_off block_size
      | _ ->
         cbc (Bytes.unsafe_to_string iv) iv_off block block_off;
         (doit [@tailcall]) block block_off block (block_off + block_size)
    in
    doit (Bytes.make block_size '\x00') 0 cbcheader 0
  in

  let rec loop ctr src src_off dst dst_off len =
    let cbcblock, cbc_off =
      match mode with
      | Encrypt -> src, src_off
      | Decrypt -> Bytes.unsafe_to_string dst, dst_off
    in
    if len = 0 then
      ()
    else if len < block_size then begin
      let buf = Bytes.make block_size '\x00' in
      Bytes.unsafe_blit dst dst_off buf 0 len ;
      ctrblock ctr buf 0 ;
      Bytes.unsafe_blit buf 0 dst dst_off len ;
      unsafe_xor_into src ~src_off dst ~dst_off len ;
      Bytes.unsafe_blit_string cbcblock cbc_off buf 0 len ;
      Bytes.unsafe_fill buf len (block_size - len) '\x00';
      cbc (Bytes.unsafe_to_string buf) 0 iv 0
    end else begin
      ctrblock ctr dst dst_off ;
      unsafe_xor_into src ~src_off dst ~dst_off block_size ;
      cbc cbcblock cbc_off iv 0 ;
      (loop [@tailcall]) (succ ctr) src (src_off + block_size) dst (dst_off + block_size) (len - block_size)
    end
  in
  loop 1 src src_off dst dst_off len;
  iv

let crypto_core ~cipher ~mode ~key ~nonce ~adata data =
  let datalen = String.length data in
  let dst = Bytes.create datalen in
  let t = crypto_core_into ~cipher ~mode ~key ~nonce ~adata data ~src_off:0 dst ~dst_off:0 datalen in
  dst, t

let crypto_t t nonce cipher key =
  let ctr = gen_ctr nonce 0 in
  cipher ~key (Bytes.unsafe_to_string ctr) ~src_off:0 ctr ~dst_off:0 ;
  unsafe_xor_into (Bytes.unsafe_to_string ctr) ~src_off:0 t ~dst_off:0 (Bytes.length t)

let unsafe_generation_encryption_into ~cipher ~key ~nonce ~adata src ~src_off dst ~dst_off ~tag_off len =
  let t = crypto_core_into ~cipher ~mode:Encrypt ~key ~nonce ~adata src ~src_off dst ~dst_off len in
  crypto_t t nonce cipher key ;
  Bytes.unsafe_blit t 0 dst tag_off block_size

let unsafe_decryption_verification_into ~cipher ~key ~nonce ~adata src ~src_off ~tag_off dst ~dst_off len =
  let tag = String.sub src tag_off block_size in
  let t = crypto_core_into ~cipher ~mode:Decrypt ~key ~nonce ~adata src ~src_off dst ~dst_off len in
  crypto_t t nonce cipher key ;
  let r  = Eqaf.equal tag (Bytes.unsafe_to_string t) in
  if not r then Bytes.unsafe_fill dst dst_off len '\000';
  r
