-- |
-- Module      : Crypto.KDF.PBKDF2
-- License     : BSD-style
-- Maintainer  : Vincent Hanquez <vincent@snarc.org>
-- Stability   : experimental
-- Portability : unknown
--
-- Password Based Key Derivation Function 2
--
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ForeignFunctionInterface #-}

module Crypto.KDF.PBKDF2
    ( PRF
    , prfHMAC
    , Parameters(..)
    , generate
    , fastPBKDF2_SHA1
    , fastPBKDF2_SHA256
    , fastPBKDF2_SHA512
    ) where

import           Data.Word
import           Data.Bits
import           Foreign.Marshal.Alloc
import           Foreign.Ptr (plusPtr, Ptr)
import           Foreign.C.Types (CUInt(..), CInt(..), CSize(..))

import           Crypto.Hash (HashAlgorithm)
import qualified Crypto.MAC.HMAC as HMAC

import           Crypto.Internal.ByteArray (ByteArray, ByteArrayAccess, Bytes)
import qualified Crypto.Internal.ByteArray as B
import           Data.Memory.PtrMethods

-- | The PRF used for PBKDF2
type PRF password =
       password -- ^ the password parameters
    -> Bytes    -- ^ the content
    -> Bytes    -- ^ prf(password,content)

-- | PRF for PBKDF2 using HMAC with the hash algorithm as parameter
prfHMAC :: (HashAlgorithm a, ByteArrayAccess password)
        => a
        -> PRF password
prfHMAC alg k = hmacIncr alg (HMAC.initialize k)
  where hmacIncr :: HashAlgorithm a => a -> HMAC.Context a -> (Bytes -> Bytes)
        hmacIncr _ !ctx = \b -> B.convert $ HMAC.finalize $ HMAC.update ctx b

-- | Parameters for PBKDF2
data Parameters = Parameters
    { iterCounts   :: Int -- ^ the number of user-defined iterations for the algorithms. e.g. WPA2 uses 4000.
    , outputLength :: Int -- ^ the number of bytes to generate out of PBKDF2
    }

-- | generate the pbkdf2 key derivation function from the output
generate :: (ByteArrayAccess password, ByteArrayAccess salt, ByteArray ba)
         => PRF password
         -> Parameters
         -> password
         -> salt
         -> ba
generate prf params password salt =
    B.allocAndFreeze (outputLength params) $ \p -> do
        memSet p 0 (outputLength params)
        loop 1 (outputLength params) p
  where
    !runPRF = prf password
    !hLen   = B.length $ runPRF B.empty

    -- run the following f function on each complete chunk.
    -- when having an incomplete chunk, we call partial.
    -- partial need to be the last call.
    --
    -- f(pass,salt,c,i) = U1 xor U2 xor .. xor Uc
    -- U1 = PRF(pass,salt || BE32(i))
    -- Uc = PRF(pass,Uc-1)
    loop iterNb len p
        | len == 0   = return ()
        | len < hLen = partial iterNb len p
        | otherwise  = do
            let applyMany 0 _     = return ()
                applyMany i uprev = do
                    let uData = runPRF uprev
                    B.withByteArray uData $ \u -> memXor p p u hLen
                    applyMany (i-1) uData
            applyMany (iterCounts params) (B.convert salt `B.append` toBS iterNb)
            loop (iterNb+1) (len - hLen) (p `plusPtr` hLen)

    partial iterNb len p = allocaBytesAligned hLen 8 $ \tmp -> do
        let applyMany :: Int -> Bytes -> IO ()
            applyMany 0 _     = return ()
            applyMany i uprev = do
                let uData = runPRF uprev
                B.withByteArray uData $ \u -> memXor tmp tmp u hLen
                applyMany (i-1) uData
        memSet tmp 0 hLen
        applyMany (iterCounts params) (B.convert salt `B.append` toBS iterNb)
        memCopy p tmp len

    -- big endian encoding of Word32
    toBS :: ByteArray ba => Word32 -> ba
    toBS w = B.pack [a,b,c,d]
      where a = fromIntegral (w `shiftR` 24)
            b = fromIntegral ((w `shiftR` 16) .&. 0xff)
            c = fromIntegral ((w `shiftR` 8) .&. 0xff)
            d = fromIntegral (w .&. 0xff)
{-# NOINLINE generate #-}

fastPBKDF2_SHA1 :: (ByteArrayAccess password, ByteArrayAccess salt, ByteArray out)
                => Parameters
                -> password
                -> salt
                -> out
fastPBKDF2_SHA1 params password salt =
    B.allocAndFreeze (outputLength params) $ \outPtr ->
    B.withByteArray password $ \passPtr ->
    B.withByteArray salt $ \saltPtr ->
        c_cryptonite_fastpbkdf2_hmac_sha1
            passPtr (fromIntegral $ B.length password)
            saltPtr (fromIntegral $ B.length salt)
            (fromIntegral $ iterCounts params)
            outPtr (fromIntegral $ outputLength params)

fastPBKDF2_SHA256 :: (ByteArrayAccess password, ByteArrayAccess salt, ByteArray out)
                  => Parameters
                  -> password
                  -> salt
                  -> out
fastPBKDF2_SHA256 params password salt =
    B.allocAndFreeze (outputLength params) $ \outPtr ->
    B.withByteArray password $ \passPtr ->
    B.withByteArray salt $ \saltPtr ->
        c_cryptonite_fastpbkdf2_hmac_sha256
            passPtr (fromIntegral $ B.length password)
            saltPtr (fromIntegral $ B.length salt)
            (fromIntegral $ iterCounts params)
            outPtr (fromIntegral $ outputLength params)

fastPBKDF2_SHA512 :: (ByteArrayAccess password, ByteArrayAccess salt, ByteArray out)
                  => Parameters
                  -> password
                  -> salt
                  -> out
fastPBKDF2_SHA512 params password salt =
    B.allocAndFreeze (outputLength params) $ \outPtr ->
    B.withByteArray password $ \passPtr ->
    B.withByteArray salt $ \saltPtr ->
        c_cryptonite_fastpbkdf2_hmac_sha512
            passPtr (fromIntegral $ B.length password)
            saltPtr (fromIntegral $ B.length salt)
            (fromIntegral $ iterCounts params)
            outPtr (fromIntegral $ outputLength params)


foreign import ccall unsafe "cryptonite_pbkdf2.h cryptonite_fastpbkdf2_hmac_sha1"
    c_cryptonite_fastpbkdf2_hmac_sha1 :: Ptr Word8 -> CSize
                                      -> Ptr Word8 -> CSize
                                      -> CUInt
                                      -> Ptr Word8 -> CSize
                                      -> IO ()

foreign import ccall unsafe "cryptonite_pbkdf2.h cryptonite_fastpbkdf2_hmac_sha256"
    c_cryptonite_fastpbkdf2_hmac_sha256 :: Ptr Word8 -> CSize
                                        -> Ptr Word8 -> CSize
                                        -> CUInt
                                        -> Ptr Word8 -> CSize
                                        -> IO ()

foreign import ccall unsafe "cryptonite_pbkdf2.h cryptonite_fastpbkdf2_hmac_sha512"
    c_cryptonite_fastpbkdf2_hmac_sha512 :: Ptr Word8 -> CSize
                                        -> Ptr Word8 -> CSize
                                        -> CUInt
                                        -> Ptr Word8 -> CSize
                                        -> IO ()