-- |

-- Module      : Crypto.Data.AFIS

-- License     : BSD-style

-- Maintainer  : Vincent Hanquez <vincent@snarc.org>

-- Stability   : experimental

-- Portability : unknown

--

-- haskell implementation of the Anti-forensic information splitter

-- available in LUKS. <http://clemens.endorphin.org/AFsplitter>

--

-- The algorithm bloats an arbitrary secret with many bits that are necessary for

-- the recovery of the key (merge), and allow greater way to permanently

-- destroy a key stored on disk.

--

{-# LANGUAGE ScopedTypeVariables #-}

module Crypto.Data.AFIS

    ( split

    , merge

    ) where

import           Crypto.Hash

import           Crypto.Random.Types

import           Crypto.Internal.Compat

import           Control.Monad (forM_, foldM)

import           Data.Word

import           Data.Bits

import           Foreign.Storable

import           Foreign.Ptr

import           Crypto.Internal.ByteArray (ByteArray, Bytes, MemView(..))

import qualified Crypto.Internal.ByteArray as B

import           Data.Memory.PtrMethods (memSet, memCopy)

-- | Split data to diffused data, using a random generator and

-- an hash algorithm.

--

-- the diffused data will consist of random data for (expandTimes-1)

-- then the last block will be xor of the accumulated random data diffused by

-- the hash algorithm.

--

-- ----------

-- -  orig  -

-- ----------

--

-- ---------- ---------- --------------

-- - rand1  - - rand2  - - orig ^ acc -

-- ---------- ---------- --------------

--

-- where acc is :

--   acc(n+1) = hash (n ++ rand(n)) ^ acc(n)

--

split :: (ByteArray ba, HashAlgorithm hash, DRG rng)

      => hash  -- ^ Hash algorithm to use as diffuser

      -> rng   -- ^ Random generator to use

      -> Int   -- ^ Number of times to diffuse the data.

      -> ba    -- ^ original data to diffuse.

      -> (ba, rng)         -- ^ The diffused data

{-# NOINLINE split #-}

split hashAlg rng expandTimes src

    | expandTimes <= 1 = error "invalid expandTimes value"

    | otherwise        = unsafeDoIO $ do

        (rng', bs) <- B.allocRet diffusedLen runOp

        return (bs, rng')

  where diffusedLen = blockSize * expandTimes

        blockSize   = B.length src

        runOp dstPtr = do

            let lastBlock = dstPtr `plusPtr` (blockSize * (expandTimes-1))

            memSet lastBlock 0 blockSize

            let randomBlockPtrs = map (plusPtr dstPtr . (*) blockSize) [0..(expandTimes-2)]

            rng' <- foldM fillRandomBlock rng randomBlockPtrs

            mapM_ (addRandomBlock lastBlock) randomBlockPtrs

            B.withByteArray src $ \srcPtr -> xorMem srcPtr lastBlock blockSize

            return rng'

        addRandomBlock lastBlock blockPtr = do

            xorMem blockPtr lastBlock blockSize

            diffuse hashAlg lastBlock blockSize

        fillRandomBlock g blockPtr = do

            let (rand :: Bytes, g') = randomBytesGenerate blockSize g

            B.withByteArray rand $ \randPtr -> memCopy blockPtr randPtr (fromIntegral blockSize)

            return g'

-- | Merge previously diffused data back to the original data.

merge :: (ByteArray ba, HashAlgorithm hash)

      => hash  -- ^ Hash algorithm used as diffuser

      -> Int   -- ^ Number of times to un-diffuse the data

      -> ba    -- ^ Diffused data

      -> ba    -- ^ Original data

{-# NOINLINE merge #-}

merge hashAlg expandTimes bs

    | r /= 0            = error "diffused data not a multiple of expandTimes"

    | originalSize <= 0 = error "diffused data null"

    | otherwise         = B.allocAndFreeze originalSize $ \dstPtr ->

        B.withByteArray bs $ \srcPtr -> do

            memSet dstPtr 0 originalSize

            forM_ [0..(expandTimes-2)] $ \i -> do

                xorMem (srcPtr `plusPtr` (i * originalSize)) dstPtr originalSize

                diffuse hashAlg dstPtr originalSize

            xorMem (srcPtr `plusPtr` ((expandTimes-1) * originalSize)) dstPtr originalSize

            return ()

  where (originalSize,r) = len `quotRem` expandTimes

        len              = B.length bs

-- | inplace Xor with an input

-- dst = src `xor` dst

xorMem :: Ptr Word8 -> Ptr Word8 -> Int -> IO ()

xorMem src dst sz

    | sz `mod` 64 == 0 = loop 8 (castPtr src :: Ptr Word64) (castPtr dst) sz

    | sz `mod` 32 == 0 = loop 4 (castPtr src :: Ptr Word32) (castPtr dst) sz

    | otherwise        = loop 1 (src :: Ptr Word8) dst sz

  where loop _    _ _ 0 = return ()

        loop incr s d n = do a <- peek s

                             b <- peek d

                             poke d (a `xor` b)

                             loop incr (s `plusPtr` incr) (d `plusPtr` incr) (n-incr)

diffuse :: HashAlgorithm hash

        => hash      -- ^ Hash function to use as diffuser

        -> Ptr Word8 -- ^ buffer to diffuse, modify in place

        -> Int       -- ^ length of buffer to diffuse

        -> IO ()

diffuse hashAlg src sz = loop src 0

  where (full,pad) = sz `quotRem` digestSize 

        loop s i

            | i < full = do h <- hashBlock i s digestSize

                            B.withByteArray h $ \hPtr -> memCopy s hPtr digestSize

                            loop (s `plusPtr` digestSize) (i+1)

            | pad /= 0 = do h <- hashBlock i s pad

                            B.withByteArray h $ \hPtr -> memCopy s hPtr pad

                            return ()

            | otherwise = return ()

        digestSize = hashDigestSize hashAlg

        -- Hash [ BE32(n), (p .. p+hashSz) ]

        hashBlock n p hashSz = do

            let ctx = hashInitWith hashAlg

            return $! hashFinalize $ hashUpdate (hashUpdate ctx (be32 n)) (MemView p hashSz)

        be32 :: Int -> Bytes

        be32 n = B.allocAndFreeze 4 $ \ptr -> do

            poke ptr               (f8 (n `shiftR` 24))

            poke (ptr `plusPtr` 1) (f8 (n `shiftR` 16))

            poke (ptr `plusPtr` 2) (f8 (n `shiftR` 8))

            poke (ptr `plusPtr` 3) (f8 n)

          where

                f8 :: Int -> Word8

                f8 = fromIntegral