-- 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 :: forall c.
BlockCipher c =>
c
-> ByteString
-> ByteString
-> ByteString
-> Either String (AuthTag, ByteString)
encryptWithOCBRFC7253 c
cipher ByteString
nonce ByteString
ad ByteString
plaintext = do
  Bool -> Either String () -> Either String ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (ByteString -> Int
B.length ByteString
nonce Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
15 Bool -> Bool -> Bool
|| ByteString -> Bool
B.null ByteString
nonce) (Either String () -> Either String ())
-> Either String () -> Either String ()
forall a b. (a -> b) -> a -> b
$
    String -> Either String ()
forall a b. a -> Either a b
Left String
"invalid nonce size for OCB"
  offset0 <- c -> ByteString -> Either String ByteString
forall c.
BlockCipher c =>
c -> ByteString -> Either String ByteString
ocbOffset0 c
cipher ByteString
nonce
  let zeroBlock = Int -> Word8 -> ByteString
B.replicate Int
16 Word8
0
      lStar = c -> ByteString -> ByteString
forall cipher ba.
(BlockCipher cipher, ByteArray ba) =>
cipher -> ba -> ba
forall ba. ByteArray ba => c -> ba -> ba
CCT.ecbEncrypt c
cipher ByteString
zeroBlock
      lDollar = ByteString -> ByteString
ocbDouble ByteString
lStar
      lCache = (ByteString -> ByteString) -> ByteString -> [ByteString]
forall a. (a -> a) -> a -> [a]
iterate ByteString -> ByteString
ocbDouble (ByteString -> ByteString
ocbDouble ByteString
lDollar)
      hashAd = c -> ByteString -> [ByteString] -> ByteString -> ByteString
forall c.
BlockCipher c =>
c -> ByteString -> [ByteString] -> ByteString -> ByteString
ocbHash c
cipher ByteString
lStar [ByteString]
lCache ByteString
ad
      (fullBlocks, partial) = splitFullAndPartial plaintext
      (cipherBlocks, offsetM, checksum) =
        foldl'
          (\([ByteString]
accBlocks, ByteString
offsetPrev, ByteString
checksumPrev) (Int
idx, ByteString
pBlock) ->
             let offsetI :: ByteString
offsetI = ByteString -> ByteString -> ByteString
xorBS ByteString
offsetPrev ([ByteString]
lCache [ByteString] -> Int -> ByteString
forall a. HasCallStack => [a] -> Int -> a
!! Int -> Int
ntz Int
idx)
                 cipherI :: ByteString
cipherI = ByteString -> ByteString -> ByteString
xorBS ByteString
offsetI (c -> ByteString -> ByteString
forall cipher ba.
(BlockCipher cipher, ByteArray ba) =>
cipher -> ba -> ba
forall ba. ByteArray ba => c -> ba -> ba
CCT.ecbEncrypt c
cipher (ByteString -> ByteString -> ByteString
xorBS ByteString
offsetI ByteString
pBlock))
                 checksumI :: ByteString
checksumI = ByteString -> ByteString -> ByteString
xorBS ByteString
checksumPrev ByteString
pBlock
              in ([ByteString]
accBlocks [ByteString] -> [ByteString] -> [ByteString]
forall a. [a] -> [a] -> [a]
++ [ByteString
cipherI], ByteString
offsetI, ByteString
checksumI))
          ([], offset0, zeroBlock)
          (zip [1 ..] fullBlocks)
      (cipherLast, offsetLast, checksumLast) =
        if B.null partial
          then (B.empty, offsetM, checksum)
          else
            let offsetStar = ByteString -> ByteString -> ByteString
xorBS ByteString
offsetM ByteString
lStar
                pad = c -> ByteString -> ByteString
forall cipher ba.
(BlockCipher cipher, ByteArray ba) =>
cipher -> ba -> ba
forall ba. ByteArray ba => c -> ba -> ba
CCT.ecbEncrypt c
cipher ByteString
offsetStar
                cipherPartial = ByteString -> ByteString -> ByteString
xorBS ByteString
partial (Int -> ByteString -> ByteString
B.take (ByteString -> Int
B.length ByteString
partial) ByteString
pad)
                checksum' = ByteString -> ByteString -> ByteString
xorBS ByteString
checksum (ByteString -> ByteString
ocbPadPartial ByteString
partial)
             in (cipherPartial, offsetStar, checksum')
      tagBytes =
        ByteString -> ByteString -> ByteString
xorBS
          (c -> ByteString -> ByteString
forall cipher ba.
(BlockCipher cipher, ByteArray ba) =>
cipher -> ba -> ba
forall ba. ByteArray ba => c -> ba -> ba
CCT.ecbEncrypt c
cipher (ByteString -> ByteString -> ByteString
xorBS (ByteString -> ByteString -> ByteString
xorBS ByteString
checksumLast ByteString
offsetLast) ByteString
lDollar))
          ByteString
hashAd
      ciphertext = [ByteString] -> ByteString
B.concat [ByteString]
cipherBlocks ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
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 :: forall c.
BlockCipher c =>
c
-> ByteString
-> ByteString
-> ByteString
-> AuthTag
-> Either String ByteString
decryptWithOCBRFC7253 =
  (ByteString
 -> ByteString
 -> ByteString
 -> ByteString
 -> ByteString
 -> ByteString
 -> String)
-> c
-> ByteString
-> ByteString
-> ByteString
-> AuthTag
-> Either String ByteString
forall c.
BlockCipher c =>
(ByteString
 -> ByteString
 -> ByteString
 -> ByteString
 -> ByteString
 -> ByteString
 -> String)
-> c
-> ByteString
-> ByteString
-> ByteString
-> AuthTag
-> Either String ByteString
decryptWithOCBRFC7253With (\ByteString
_ ByteString
_ ByteString
_ ByteString
_ ByteString
_ ByteString
_ -> String
"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 :: forall c.
BlockCipher c =>
(ByteString
 -> ByteString
 -> ByteString
 -> ByteString
 -> ByteString
 -> ByteString
 -> String)
-> c
-> ByteString
-> ByteString
-> ByteString
-> AuthTag
-> Either String ByteString
decryptWithOCBRFC7253With ByteString
-> ByteString
-> ByteString
-> ByteString
-> ByteString
-> ByteString
-> String
onAuthFailure c
cipher ByteString
nonce ByteString
ad ByteString
ciphertext AuthTag
authTag = do
  Bool -> Either String () -> Either String ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (ByteString -> Int
B.length ByteString
nonce Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
15 Bool -> Bool -> Bool
|| ByteString -> Bool
B.null ByteString
nonce) (Either String () -> Either String ())
-> Either String () -> Either String ()
forall a b. (a -> b) -> a -> b
$
    String -> Either String ()
forall a b. a -> Either a b
Left String
"invalid nonce size for OCB"
  Bool -> Either String () -> Either String ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (ByteString -> Int
B.length ByteString
tagBytes Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
16) (Either String () -> Either String ())
-> Either String () -> Either String ()
forall a b. (a -> b) -> a -> b
$
    String -> Either String ()
forall a b. a -> Either a b
Left String
"invalid auth tag size for OCB"
  offset0 <- c -> ByteString -> Either String ByteString
forall c.
BlockCipher c =>
c -> ByteString -> Either String ByteString
ocbOffset0 c
cipher ByteString
nonce
  let zeroBlock = Int -> Word8 -> ByteString
B.replicate Int
16 Word8
0
      lStar = c -> ByteString -> ByteString
forall cipher ba.
(BlockCipher cipher, ByteArray ba) =>
cipher -> ba -> ba
forall ba. ByteArray ba => c -> ba -> ba
CCT.ecbEncrypt c
cipher ByteString
zeroBlock
      lDollar = ByteString -> ByteString
ocbDouble ByteString
lStar
      lCache = (ByteString -> ByteString) -> ByteString -> [ByteString]
forall a. (a -> a) -> a -> [a]
iterate ByteString -> ByteString
ocbDouble (ByteString -> ByteString
ocbDouble ByteString
lDollar)
      hashAd = c -> ByteString -> [ByteString] -> ByteString -> ByteString
forall c.
BlockCipher c =>
c -> ByteString -> [ByteString] -> ByteString -> ByteString
ocbHash c
cipher ByteString
lStar [ByteString]
lCache ByteString
ad
      (fullBlocks, partial) = splitFullAndPartial ciphertext
      (plainBlocks, offsetM, checksum) =
        foldl'
          (\([ByteString]
accBlocks, ByteString
offsetPrev, ByteString
checksumPrev) (Int
idx, ByteString
cBlock) ->
             let offsetI :: ByteString
offsetI = ByteString -> ByteString -> ByteString
xorBS ByteString
offsetPrev ([ByteString]
lCache [ByteString] -> Int -> ByteString
forall a. HasCallStack => [a] -> Int -> a
!! Int -> Int
ntz Int
idx)
                 plainI :: ByteString
plainI = ByteString -> ByteString -> ByteString
xorBS ByteString
offsetI (c -> ByteString -> ByteString
forall cipher ba.
(BlockCipher cipher, ByteArray ba) =>
cipher -> ba -> ba
forall ba. ByteArray ba => c -> ba -> ba
CCT.ecbDecrypt c
cipher (ByteString -> ByteString -> ByteString
xorBS ByteString
offsetI ByteString
cBlock))
                 checksumI :: ByteString
checksumI = ByteString -> ByteString -> ByteString
xorBS ByteString
checksumPrev ByteString
plainI
              in ([ByteString]
accBlocks [ByteString] -> [ByteString] -> [ByteString]
forall a. [a] -> [a] -> [a]
++ [ByteString
plainI], ByteString
offsetI, ByteString
checksumI))
          ([], offset0, zeroBlock)
          (zip [1 ..] fullBlocks)
      (plainLast, offsetLast, checksumLast) =
        if B.null partial
          then (B.empty, offsetM, checksum)
          else
            let offsetStar = ByteString -> ByteString -> ByteString
xorBS ByteString
offsetM ByteString
lStar
                pad = c -> ByteString -> ByteString
forall cipher ba.
(BlockCipher cipher, ByteArray ba) =>
cipher -> ba -> ba
forall ba. ByteArray ba => c -> ba -> ba
CCT.ecbEncrypt c
cipher ByteString
offsetStar
                plainPartial = ByteString -> ByteString -> ByteString
xorBS ByteString
partial (Int -> ByteString -> ByteString
B.take (ByteString -> Int
B.length ByteString
partial) ByteString
pad)
                checksum' = ByteString -> ByteString -> ByteString
xorBS ByteString
checksum (ByteString -> ByteString
ocbPadPartial ByteString
plainPartial)
             in (plainPartial, offsetStar, checksum')
      tagComputed =
        ByteString -> ByteString -> ByteString
xorBS
          (c -> ByteString -> ByteString
forall cipher ba.
(BlockCipher cipher, ByteArray ba) =>
cipher -> ba -> ba
forall ba. ByteArray ba => c -> ba -> ba
CCT.ecbEncrypt c
cipher (ByteString -> ByteString -> ByteString
xorBS (ByteString -> ByteString -> ByteString
xorBS ByteString
checksumLast ByteString
offsetLast) ByteString
lDollar))
          ByteString
hashAd
      plaintext = [ByteString] -> ByteString
B.concat [ByteString]
plainBlocks ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
plainLast
      computedTag = Int -> ByteString -> ByteString
B.take Int
16 ByteString
tagComputed
   in if BA.constEq tagBytes computedTag
        then Right plaintext
        else Left (onAuthFailure tagBytes computedTag nonce ad hashAd plaintext)
  where
    tagBytes :: ByteString
tagBytes = AuthTag -> ByteString
forall bin bout.
(ByteArrayAccess bin, ByteArray bout) =>
bin -> bout
BA.convert AuthTag
authTag :: B.ByteString

mkAuthTag :: B.ByteString -> CCT.AuthTag
mkAuthTag :: ByteString -> AuthTag
mkAuthTag = Bytes -> AuthTag
CCT.AuthTag (Bytes -> AuthTag)
-> (ByteString -> Bytes) -> ByteString -> AuthTag
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Bytes
forall bin bout.
(ByteArrayAccess bin, ByteArray bout) =>
bin -> bout
BA.convert

ocbOffset0 :: CCT.BlockCipher c => c -> B.ByteString -> Either String B.ByteString
ocbOffset0 :: forall c.
BlockCipher c =>
c -> ByteString -> Either String ByteString
ocbOffset0 c
cipher ByteString
nonce = do
  let nonceLen :: Int
nonceLen = ByteString -> Int
B.length ByteString
nonce
      prefixLen :: Int
prefixLen = Int
16 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
nonceLen
  Bool -> Either String () -> Either String ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
prefixLen Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0) (Either String () -> Either String ())
-> Either String () -> Either String ()
forall a b. (a -> b) -> a -> b
$
    String -> Either String ()
forall a b. a -> Either a b
Left String
"invalid nonce size for OCB"
  let prefix :: ByteString
prefix = [Word8] -> ByteString
B.pack (Int -> Word8 -> [Word8]
forall a. Int -> a -> [a]
replicate (Int
prefixLen Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Word8
0 [Word8] -> [Word8] -> [Word8]
forall a. Semigroup a => a -> a -> a
<> [Word8
1 :: Word8])
      nonceBlock :: ByteString
nonceBlock = ByteString
prefix ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
nonce
      bottom :: Int
bottom = Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (HasCallStack => ByteString -> Word8
ByteString -> Word8
B.last ByteString
nonceBlock Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x3f) :: Int
      nonceTop :: ByteString
nonceTop = HasCallStack => ByteString -> ByteString
ByteString -> ByteString
B.init ByteString
nonceBlock ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> Word8 -> ByteString
B.singleton (HasCallStack => ByteString -> Word8
ByteString -> Word8
B.last ByteString
nonceBlock Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0xc0)
      kTop :: ByteString
kTop = c -> ByteString -> ByteString
forall cipher ba.
(BlockCipher cipher, ByteArray ba) =>
cipher -> ba -> ba
forall ba. ByteArray ba => c -> ba -> ba
CCT.ecbEncrypt c
cipher ByteString
nonceTop
      stretch :: ByteString
stretch = ByteString
kTop ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString -> ByteString -> ByteString
xorBS (Int -> ByteString -> ByteString
B.take Int
8 ByteString
kTop) (Int -> ByteString -> ByteString
B.take Int
8 (Int -> ByteString -> ByteString
B.drop Int
1 ByteString
kTop))
  ByteString -> Either String ByteString
forall a b. b -> Either a b
Right (ByteString -> Int -> ByteString
ocbBitSlice128 ByteString
stretch Int
bottom)

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

ocbPadPartial :: B.ByteString -> B.ByteString
ocbPadPartial :: ByteString -> ByteString
ocbPadPartial ByteString
bs = ByteString
bs ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> Word8 -> ByteString
B.singleton Word8
0x80 ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> Int -> Word8 -> ByteString
B.replicate (Int
15 Int -> Int -> Int
forall a. Num a => a -> a -> a
- ByteString -> Int
B.length ByteString
bs) Word8
0

ocbDouble :: B.ByteString -> B.ByteString
ocbDouble :: ByteString -> ByteString
ocbDouble ByteString
bs =
  let shifted :: ByteString
shifted = ByteString -> ByteString
shiftLeftOne ByteString
bs
      carry :: Bool
carry = (HasCallStack => ByteString -> Word8
ByteString -> Word8
B.head ByteString
bs Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x80) Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word8
0
   in if Bool
carry
        then HasCallStack => ByteString -> ByteString
ByteString -> ByteString
B.init ByteString
shifted ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> Word8 -> ByteString
B.singleton (HasCallStack => ByteString -> Word8
ByteString -> Word8
B.last ByteString
shifted Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
`xor` Word8
0x87)
        else ByteString
shifted

shiftLeftOne :: B.ByteString -> B.ByteString
shiftLeftOne :: ByteString -> ByteString
shiftLeftOne ByteString
bs = [Word8] -> ByteString
B.pack [Word8]
shifted
  where
    ([Word8]
shifted, Word8
_) =
      (([Word8], Word8) -> Word8 -> ([Word8], Word8))
-> ([Word8], Word8) -> [Word8] -> ([Word8], Word8)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl'
        (\([Word8]
acc, Word8
carryIn) Word8
x ->
           let y :: Word8
y = ((Word8
x Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
1) Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0xff) Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Word8
carryIn
               carryOut :: Word8
carryOut = if (Word8
x Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x80) Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word8
0 then Word8
1 else Word8
0
            in (Word8
y Word8 -> [Word8] -> [Word8]
forall a. a -> [a] -> [a]
: [Word8]
acc, Word8
carryOut))
        ([], Word8
0)
        ([Word8] -> [Word8]
forall a. [a] -> [a]
reverse (ByteString -> [Word8]
B.unpack ByteString
bs))

xorBS :: B.ByteString -> B.ByteString -> B.ByteString
xorBS :: ByteString -> ByteString -> ByteString
xorBS ByteString
a ByteString
b = [Word8] -> ByteString
B.pack ((Word8 -> Word8 -> Word8) -> ByteString -> ByteString -> [Word8]
forall a. (Word8 -> Word8 -> a) -> ByteString -> ByteString -> [a]
B.zipWith Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
xor ByteString
a ByteString
b)

ocbBitSlice128 :: B.ByteString -> Int -> B.ByteString
ocbBitSlice128 :: ByteString -> Int -> ByteString
ocbBitSlice128 ByteString
stretch Int
startBit =
  let stretchInt :: Integer
stretchInt = ByteString -> Integer
forall ba. ByteArrayAccess ba => ba -> Integer
os2ip ByteString
stretch
      shift :: Int
shift = Int
192 Int -> Int -> Int
forall a. Num a => a -> a -> a
- (Int
startBit Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
128)
      mask :: Integer
mask = (Integer
1 Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
`shiftL` (Int
128 :: Int)) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1
      slice :: Integer
slice = (Integer
stretchInt Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
`shiftR` Int
shift) Integer -> Integer -> Integer
forall a. Bits a => a -> a -> a
.&. Integer
mask
   in ByteString -> ByteString
leftPadTo16 (Integer -> ByteString
forall ba. ByteArray ba => Integer -> ba
i2osp Integer
slice)

leftPadTo16 :: B.ByteString -> B.ByteString
leftPadTo16 :: ByteString -> ByteString
leftPadTo16 ByteString
bs
  | ByteString -> Int
B.length ByteString
bs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
16 = Int -> ByteString -> ByteString
B.drop (ByteString -> Int
B.length ByteString
bs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
16) ByteString
bs
  | Bool
otherwise = Int -> Word8 -> ByteString
B.replicate (Int
16 Int -> Int -> Int
forall a. Num a => a -> a -> a
- ByteString -> Int
B.length ByteString
bs) Word8
0 ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
bs

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

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

ntz :: Int -> Int
ntz :: Int -> Int
ntz Int
i = Int -> Int -> Int
forall {t} {t}. (Bits t, Num t, Num t) => t -> t -> t
go Int
i Int
0
  where
    go :: t -> t -> t
go t
n t
c
      | t
n t -> t -> t
forall a. Bits a => a -> a -> a
.&. t
1 t -> t -> Bool
forall a. Eq a => a -> a -> Bool
== t
1 = t
c
      | Bool
otherwise = t -> t -> t
go (t
n t -> Int -> t
forall a. Bits a => a -> Int -> a
`shiftR` Int
1) (t
c t -> t -> t
forall a. Num a => a -> a -> a
+ t
1)