## Nim-Codex
## Copyright (c) 2025 Status Research & Development GmbH
## Licensed under either of
##  * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE))
##  * MIT license ([LICENSE-MIT](LICENSE-MIT))
## at your option.
## This file may not be copied, modified, or distributed except according to
## those terms.

import std/sequtils

import pkg/bearssl/[blockx, hash]
import pkg/stew/[byteutils, endians2]
import pkg/questionable/results

import ../rng
import ./bearsslhash

{.push raises: [].}

const
  MasterKeySize = 32 # 256 bits
  KeySize = 24 # 192 bits for AES-192
  IvSize = 16 # 128 bits
  KeyDerivationIdentifier = "aes192_block_key".toBytes
  IvDerivationIdentifier = "aes192_block_iv".toBytes

type CodexEncryption* = ref object
  masterKey: seq[byte]

proc newCodexEncryption*(): CodexEncryption =
  let masterKey = newSeqWith(MasterKeySize, Rng.instance.rand(uint8.high).byte)
  CodexEncryption(masterKey: masterKey)

proc newCodexEncryption*(masterKey: seq[byte]): CodexEncryption =
  CodexEncryption(masterKey: masterKey)

proc getKeyHexEncoded*(self: CodexEncryption): string =
  self.masterKey.toHex()

proc deriveKeyForBlockIndex(self: CodexEncryption, blockIndex: uint32): seq[byte] =
  let blockIndexArray = toBytes(blockIndex, bigEndian)
  bearSslHash(
    addr sha256Vtable, self.masterKey & KeyDerivationIdentifier & blockIndexArray.toSeq
  )[0 ..< KeySize]

proc deriveIvForBlockIndex(self: CodexEncryption, blockIndex: uint32): seq[byte] =
  let blockIndexArray = toBytes(blockIndex, bigEndian)
  bearSslHash(
    addr sha256Vtable, self.masterKey & IvDerivationIdentifier & blockIndexArray.toSeq
  )[0 ..< IvSize]

proc encryptBlock*(
    self: CodexEncryption, blockData: seq[byte], blockIndex: uint32
): ?!seq[byte] =
  if blockData.len mod IvSize != 0:
    return failure("Block data length is not a multiple of 16!")

  let key = self.deriveKeyForBlockIndex(blockIndex)
  let iv = self.deriveIvForBlockIndex(blockIndex)
  var outBuff = newSeqUninit[byte](blockData.len)
  copyMem(addr outBuff[0], addr blockData[0], blockData.len)

  var encCtx: AesBigCbcencKeys
  aesBigCbcencInit(encCtx, addr key[0], key.len.uint)
  aesBigCbcencRun(encCtx, addr iv[0], addr outBuff[0], outBuff.len.uint)
  success move outBuff

proc decryptBlock*(
    self: CodexEncryption, encryptedBlockData: seq[byte], blockIndex: uint32
): ?!seq[byte] =
  if encryptedBlockData.len mod IvSize != 0:
    return failure("Encrypted block data length is not a multiple of 16!")

  let key = self.deriveKeyForBlockIndex(blockIndex)
  let iv = self.deriveIvForBlockIndex(blockIndex)
  var outBuff = newSeqUninit[byte](encryptedBlockData.len)
  copyMem(addr outBuff[0], addr encryptedBlockData[0], encryptedBlockData.len)

  var decCtx: AesBigCbcdecKeys
  aesBigCbcdecInit(decCtx, addr key[0], key.len.uint)
  aesBigCbcdecRun(decCtx, addr iv[0], addr outBuff[0], outBuff.len.uint)
  success move outBuff