let nonce_size = 12
let tag_size = 16
let version_size = 1
let encryption_version = 1
let key_size = 16
let overhead = version_size + nonce_size + tag_size

type key = Mirage_crypto.AES.GCM.key

let init_key secret =
  if String.length secret <> key_size then Error `Invalid_key_length
  else Ok (Mirage_crypto.AES.GCM.of_secret secret)

let generate_nonce (random : _ Eio.Flow.source) =
  let buf = Cstruct.create nonce_size in
  Eio.Flow.read_exact random buf;
  Cstruct.to_string buf

let encrypt ~key ~random plaintext =
  let nonce = generate_nonce random in
  let ciphertext =
    Mirage_crypto.AES.GCM.authenticate_encrypt ~key ~nonce
      (Cstruct.to_string plaintext)
  in
  let result =
    Cstruct.create (version_size + nonce_size + String.length ciphertext)
  in
  Cstruct.set_uint8 result 0 encryption_version;
  Cstruct.blit_from_string nonce 0 result version_size nonce_size;
  Cstruct.blit_from_string ciphertext 0 result
    (version_size + nonce_size)
    (String.length ciphertext);
  result

let pkcs7_unpad data block_size =
  let len = String.length data in
  if len = 0 then data
  else
    let padding = Char.code data.[len - 1] in
    if padding > 0 && padding <= block_size && padding <= len then
      String.sub data 0 (len - padding)
    else data

let decrypt ~key data =
  if Cstruct.length data < overhead then Error `Too_short
  else
    let version = Cstruct.get_uint8 data 0 in
    if version > 1 then Error `Unsupported_version
    else
      let nonce =
        Cstruct.to_string (Cstruct.sub data version_size nonce_size)
      in
      let ciphertext =
        Cstruct.to_string
          (Cstruct.sub data
             (version_size + nonce_size)
             (Cstruct.length data - version_size - nonce_size))
      in
      match
        Mirage_crypto.AES.GCM.authenticate_decrypt ~key ~nonce ciphertext
      with
      | Some plaintext ->
          let plaintext =
            if version = 0 then pkcs7_unpad plaintext 16 else plaintext
          in
          Ok (Cstruct.of_string plaintext)
      | None -> Error `Decryption_failed