-- RFC7253OCB.hs: RFC7253 implementation as a
--                workaround until crypton is fixed
-- Copyright © 2026  Clint Adams
-- This software is released under the terms of the Expat license.
-- (See the LICENSE file).
{-# LANGUAGE PackageImports #-}

module Codec.Encryption.OpenPGP.Internal.RFC7253OCB
  ( encryptWithOCBRFC7253
  , decryptWithOCBRFC7253
  , decryptWithOCBRFC7253With
  ) where

import Control.Monad (when)
import qualified "crypton" Crypto.Cipher.Types as CCT
import Data.Bits ((.&.), (.|.), shiftL, shiftR, xor)
import qualified Data.ByteArray as BA
import qualified Data.ByteString as B
import Data.List (foldl')
import Crypto.Number.Serialize (i2osp, os2ip)
import Data.Word (Word8)

encryptWithOCBRFC7253 ::
     CCT.BlockCipher c
  => c
  -> B.ByteString
  -> B.ByteString
  -> B.ByteString
  -> Either String (CCT.AuthTag, B.ByteString)
encryptWithOCBRFC7253 cipher nonce ad plaintext = do
  when (B.length nonce > 15 || B.null nonce) $
    Left "invalid nonce size for OCB"
  offset0 <- ocbOffset0 cipher nonce
  let zeroBlock = B.replicate 16 0
      lStar = CCT.ecbEncrypt cipher zeroBlock
      lDollar = ocbDouble lStar
      lCache = iterate ocbDouble (ocbDouble lDollar)
      hashAd = ocbHash cipher lStar lCache ad
      (fullBlocks, partial) = splitFullAndPartial plaintext
      (cipherBlocks, offsetM, checksum) =
        foldl'
          (\(accBlocks, offsetPrev, checksumPrev) (idx, pBlock) ->
             let offsetI = xorBS offsetPrev (lCache !! ntz idx)
                 cipherI = xorBS offsetI (CCT.ecbEncrypt cipher (xorBS offsetI pBlock))
                 checksumI = xorBS checksumPrev pBlock
              in (accBlocks ++ [cipherI], offsetI, checksumI))
          ([], offset0, zeroBlock)
          (zip [1 ..] fullBlocks)
      (cipherLast, offsetLast, checksumLast) =
        if B.null partial
          then (B.empty, offsetM, checksum)
          else
            let offsetStar = xorBS offsetM lStar
                pad = CCT.ecbEncrypt cipher offsetStar
                cipherPartial = xorBS partial (B.take (B.length partial) pad)
                checksum' = xorBS checksum (ocbPadPartial partial)
             in (cipherPartial, offsetStar, checksum')
      tagBytes =
        xorBS
          (CCT.ecbEncrypt cipher (xorBS (xorBS checksumLast offsetLast) lDollar))
          hashAd
      ciphertext = B.concat cipherBlocks <> cipherLast
   in Right (mkAuthTag (B.take 16 tagBytes), ciphertext)

decryptWithOCBRFC7253 ::
     CCT.BlockCipher c
  => c
  -> B.ByteString
  -> B.ByteString
  -> B.ByteString
  -> CCT.AuthTag
  -> Either String B.ByteString
decryptWithOCBRFC7253 =
  decryptWithOCBRFC7253With (\_ _ _ _ _ _ -> "OCB authentication failed")

decryptWithOCBRFC7253With ::
     CCT.BlockCipher c
  => (B.ByteString -> B.ByteString -> B.ByteString -> B.ByteString -> B.ByteString -> B.ByteString -> String)
  -> c
  -> B.ByteString
  -> B.ByteString
  -> B.ByteString
  -> CCT.AuthTag
  -> Either String B.ByteString
decryptWithOCBRFC7253With onAuthFailure cipher nonce ad ciphertext authTag = do
  when (B.length nonce > 15 || B.null nonce) $
    Left "invalid nonce size for OCB"
  when (B.length tagBytes /= 16) $
    Left "invalid auth tag size for OCB"
  offset0 <- ocbOffset0 cipher nonce
  let zeroBlock = B.replicate 16 0
      lStar = CCT.ecbEncrypt cipher zeroBlock
      lDollar = ocbDouble lStar
      lCache = iterate ocbDouble (ocbDouble lDollar)
      hashAd = ocbHash cipher lStar lCache ad
      (fullBlocks, partial) = splitFullAndPartial ciphertext
      (plainBlocks, offsetM, checksum) =
        foldl'
          (\(accBlocks, offsetPrev, checksumPrev) (idx, cBlock) ->
             let offsetI = xorBS offsetPrev (lCache !! ntz idx)
                 plainI = xorBS offsetI (CCT.ecbDecrypt cipher (xorBS offsetI cBlock))
                 checksumI = xorBS checksumPrev plainI
              in (accBlocks ++ [plainI], offsetI, checksumI))
          ([], offset0, zeroBlock)
          (zip [1 ..] fullBlocks)
      (plainLast, offsetLast, checksumLast) =
        if B.null partial
          then (B.empty, offsetM, checksum)
          else
            let offsetStar = xorBS offsetM lStar
                pad = CCT.ecbEncrypt cipher offsetStar
                plainPartial = xorBS partial (B.take (B.length partial) pad)
                checksum' = xorBS checksum (ocbPadPartial plainPartial)
             in (plainPartial, offsetStar, checksum')
      tagComputed =
        xorBS
          (CCT.ecbEncrypt cipher (xorBS (xorBS checksumLast offsetLast) lDollar))
          hashAd
      plaintext = B.concat plainBlocks <> plainLast
      computedTag = B.take 16 tagComputed
   in if BA.constEq tagBytes computedTag
        then Right plaintext
        else Left (onAuthFailure tagBytes computedTag nonce ad hashAd plaintext)
  where
    tagBytes = BA.convert authTag :: B.ByteString

mkAuthTag :: B.ByteString -> CCT.AuthTag
mkAuthTag = CCT.AuthTag . BA.convert

ocbOffset0 :: CCT.BlockCipher c => c -> B.ByteString -> Either String B.ByteString
ocbOffset0 cipher nonce = do
  let nonceLen = B.length nonce
      prefixLen = 16 - nonceLen
  when (prefixLen <= 0) $
    Left "invalid nonce size for OCB"
  let prefix = B.pack (replicate (prefixLen - 1) 0 <> [1 :: Word8])
      nonceBlock = prefix <> nonce
      bottom = fromIntegral (B.last nonceBlock .&. 0x3f) :: Int
      nonceTop = B.init nonceBlock <> B.singleton (B.last nonceBlock .&. 0xc0)
      kTop = CCT.ecbEncrypt cipher nonceTop
      stretch = kTop <> xorBS (B.take 8 kTop) (B.take 8 (B.drop 1 kTop))
  Right (ocbBitSlice128 stretch bottom)

ocbHash ::
     CCT.BlockCipher c
  => c
  -> B.ByteString
  -> [B.ByteString]
  -> B.ByteString
  -> B.ByteString
ocbHash cipher lStar lCache ad =
  let (fullBlocks, partial) = splitFullAndPartial ad
      (sumBlocks, offsetFinal) =
        foldl'
          (\(acc, offsetPrev) (idx, block) ->
             let offsetI = xorBS offsetPrev (lCache !! ntz idx)
                 sumI = xorBS acc (CCT.ecbEncrypt cipher (xorBS offsetI block))
              in (sumI, offsetI))
          (B.replicate 16 0, B.replicate 16 0)
          (zip [1 ..] fullBlocks)
   in if B.null partial
        then sumBlocks
        else
          let offsetStar = xorBS offsetFinal lStar
              block = ocbPadPartial partial
           in xorBS sumBlocks (CCT.ecbEncrypt cipher (xorBS offsetStar block))

ocbPadPartial :: B.ByteString -> B.ByteString
ocbPadPartial bs = bs <> B.singleton 0x80 <> B.replicate (15 - B.length bs) 0

ocbDouble :: B.ByteString -> B.ByteString
ocbDouble bs =
  let shifted = shiftLeftOne bs
      carry = (B.head bs .&. 0x80) /= 0
   in if carry
        then B.init shifted <> B.singleton (B.last shifted `xor` 0x87)
        else shifted

shiftLeftOne :: B.ByteString -> B.ByteString
shiftLeftOne bs = B.pack shifted
  where
    (shifted, _) =
      foldl'
        (\(acc, carryIn) x ->
           let y = ((x `shiftL` 1) .&. 0xff) .|. carryIn
               carryOut = if (x .&. 0x80) /= 0 then 1 else 0
            in (y : acc, carryOut))
        ([], 0)
        (reverse (B.unpack bs))

xorBS :: B.ByteString -> B.ByteString -> B.ByteString
xorBS a b = B.pack (B.zipWith xor a b)

ocbBitSlice128 :: B.ByteString -> Int -> B.ByteString
ocbBitSlice128 stretch startBit =
  let stretchInt = os2ip stretch
      shift = 192 - (startBit + 128)
      mask = (1 `shiftL` (128 :: Int)) - 1
      slice = (stretchInt `shiftR` shift) .&. mask
   in leftPadTo16 (i2osp slice)

leftPadTo16 :: B.ByteString -> B.ByteString
leftPadTo16 bs
  | B.length bs >= 16 = B.drop (B.length bs - 16) bs
  | otherwise = B.replicate (16 - B.length bs) 0 <> bs

splitFullAndPartial :: B.ByteString -> ([B.ByteString], B.ByteString)
splitFullAndPartial bs
  | B.null bs = ([], B.empty)
  | otherwise =
      let fullLen = (B.length bs `div` 16) * 16
          (fullPart, rest) = B.splitAt fullLen bs
       in (chunk16 fullPart, rest)

chunk16 :: B.ByteString -> [B.ByteString]
chunk16 bs
  | B.null bs = []
  | otherwise =
      let (h, t) = B.splitAt 16 bs
       in h : chunk16 t

ntz :: Int -> Int
ntz i = go i 0
  where
    go n c
      | n .&. 1 == 1 = c
      | otherwise = go (n `shiftR` 1) (c + 1)
