{-# LANGUAGE ExistentialQuantification, MultiParamTypeClasses #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Crypto.Nettle.Hash.Types
-- Copyright   :  (c) 2013 Stefan Bühler
-- License     :  MIT-style (see the file COPYING)
-- 
-- Maintainer  :  stbuehler@web.de
-- Stability   :  experimental
-- Portability :  portable
--
-- Collection of internal types due to cyclic dependencies
--
-----------------------------------------------------------------------------

module Crypto.Nettle.Hash.Types
	( HashAlgorithm(..)
	, hash
	, hash'
	, hashLazy
	, hashLazy'

	, KeyedHashAlgorithm(..)
	, KeyedHash(..)

	, keyedHashDigestSize
	, keyedHashDigestSize'
	, keyedHashName
	, keyedHashName'
	, keyedHashInit
	, keyedHashInit'
	, keyedHashUpdate
	, keyedHashUpdateLazy
	, keyedHashFinalize
	, keyedHash
	, keyedHash'
	, keyedHashLazy
	, keyedHashLazy'

	, module Data.Tagged

	, HMAC
	, hmacInit
	, hmacInit'
	, hmac
	, hmac'
	, hmacLazy
	, hmacLazy'
	) where

import Data.Tagged
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as L
import Control.Applicative ((<$>))
import Data.Bits (xor)
import Data.List (foldl')


{-|
'HashAlgorithm' is a class that hash algorithms will implement. generating a digest is a 3 step procedure:

  * 'hashInit' to create a new context

  * 'hashUpdate' to hash data

  * 'hashFinalize' to extract the final digest

The final digest has 'hashDigestSize' bytes, and the algorithm uses 'hashBlockSize' as internal block size.
-}
class HashAlgorithm a where
	-- | Block size in bytes the hash algorithm operates on
	hashBlockSize  :: Tagged a Int
	-- | Digest size in bytes the hash algorithm returns
	hashDigestSize :: Tagged a Int
	-- | Name of the hash algorithm
	hashName       :: Tagged a String
	-- | Initialize a new context for this hash algorithm
	hashInit       :: a
	-- | Update the context with bytestring, and return a new context with the updates.
	hashUpdate     :: a -> B.ByteString -> a
	-- | Update the context with a lazy bytestring, and return a new context with the updates.
	hashUpdateLazy :: a -> L.ByteString -> a
	hashUpdateLazy a = foldl' hashUpdate a . L.toChunks
	-- | Finalize a context and return a digest.
	hashFinalize   :: a -> B.ByteString
	-- | Use 'HashAlgorithm' for HMAC; can use a optimized variant or the default 'hmacInit' one
	hashHMAC       :: B.ByteString -> Tagged a KeyedHash
	hashHMAC = hmacInit

{-|
Helper to hash a single (strict) 'B.ByteString' in one step.

Example:

> untag (hash (fromString "abc") :: Tagged SHA256 B.ByteString)
-}
hash :: HashAlgorithm a => B.ByteString -> Tagged a B.ByteString
hash msg = hashFinalize <$> flip hashUpdate msg <$> tagSelf hashInit
{-|
Untagged variant of 'hash'; takes a (possible 'undefined') typed 'HashAlgorithm' context as parameter.

Example:

> hash' (undefined :: SHA256) $ fromString "abc"
-}
hash' :: HashAlgorithm a => a -> B.ByteString -> B.ByteString
hash' a = flip witness a . hash

{-|
Helper to hash a single (lazy) 'L.ByteString' in one step.

Example:

> untag (hashLazy (fromString "abc") :: Tagged SHA256 L.ByteString)
-}
hashLazy :: HashAlgorithm a => L.ByteString -> Tagged a L.ByteString
hashLazy msg = L.fromStrict <$> hashFinalize <$> flip hashUpdateLazy msg <$> tagSelf hashInit
{-|
Untagged variant of 'hashLazy'; takes a (possible 'undefined') typed 'HashAlgorithm' context as parameter.

Example:

> hashLazy' (undefined :: SHA256) $ fromString "abc"
-}
hashLazy' :: HashAlgorithm a => a -> L.ByteString -> L.ByteString
hashLazy' a = flip witness a . hashLazy

{-|
'KeyedHashAlgorithm' is a class for keyed hash algorithms that take a key and a message to produce a digest.
The most popular example is 'HMAC'.
-}
class KeyedHashAlgorithm k where
	-- | Digest size in bytes the keyed hash algorithm returns
	implKeyedHashDigestSize :: Tagged k Int
	-- | Name
	implKeyedHashName :: Tagged k String
	-- | Initialize state from a key
	implKeyedHashInit :: B.ByteString -> k
	-- | Add more message data to the state
	implKeyedHashUpdate :: k -> B.ByteString -> k
	-- | Add more lazy message data to the state
	implKeyedHashUpdateLazy :: k -> L.ByteString -> k
	implKeyedHashUpdateLazy k = foldl' implKeyedHashUpdate k . L.toChunks
	-- | Produce final digest
	implKeyedHashFinalize :: k -> B.ByteString

{-|
'KeyedHash' hides the 'KeyedHashAlgorithm' implementation.
-}
data KeyedHash = forall k. KeyedHashAlgorithm k => KeyedHash !k

{-|
Untagged variant of 'implKeyedHashDigestSize'; takes a (possible 'undefined') key typed value from a 'KeyedHashAlgorithm' instance as parameter.
-}
keyedHashDigestSize :: KeyedHashAlgorithm k => k -> Int
keyedHashDigestSize k = implKeyedHashDigestSize `witness` k
{-|
Get 'implKeyedHashDigestSize' from a 'KeyedHash'
-}
keyedHashDigestSize' :: KeyedHash -> Int
keyedHashDigestSize' (KeyedHash k) = implKeyedHashDigestSize `witness` k
{-|
Untagged variant of 'implKeyedHashName'; takes a (possible 'undefined') key typed value from a 'KeyedHashAlgorithm' instance as parameter.
-}
keyedHashName :: KeyedHashAlgorithm k => k -> String
keyedHashName k = implKeyedHashName `witness` k
{-|
Get 'implKeyedHashName' from a 'KeyedHash'
-}
keyedHashName' :: KeyedHash -> String
keyedHashName' (KeyedHash k) = implKeyedHashName `witness` k
{-|
Initialize a 'KeyedHash' context from a @key@
-}
keyedHashInit :: KeyedHashAlgorithm k => B.ByteString {- ^ @key@ argument -} -> Tagged k KeyedHash
keyedHashInit key = KeyedHash <$> tagSelf (implKeyedHashInit key)
{-|
Untagged variant of 'keyedHashInit'; takes a (possible 'undefined') key typed value from a 'KeyedHashAlgorithm' instance as parameter.
-}
keyedHashInit' :: KeyedHashAlgorithm k => k -> B.ByteString -> KeyedHash
keyedHashInit' k key = keyedHashInit key `witness` k
{-|
Add more message data to the context
-}
keyedHashUpdate :: KeyedHash -> B.ByteString -> KeyedHash
keyedHashUpdate (KeyedHash k) = KeyedHash . implKeyedHashUpdate k
{-|
Add more lazy message data to the context
-}
keyedHashUpdateLazy :: KeyedHash -> L.ByteString -> KeyedHash
keyedHashUpdateLazy (KeyedHash k) = KeyedHash . implKeyedHashUpdateLazy k
{-|
Produce final digest
-}
keyedHashFinalize :: KeyedHash -> B.ByteString
keyedHashFinalize (KeyedHash k) = implKeyedHashFinalize k
{-|
Helper to hash @key@ and @message@ in one step

Example:

> untag (keyedHash (fromString "secretkey") (fromString "secret message") :: Tagged (HMAC SHA256) B.ByteString)
-}
keyedHash :: KeyedHashAlgorithm k => B.ByteString -> B.ByteString -> Tagged k B.ByteString
keyedHash key msg = keyedHashFinalize <$> flip keyedHashUpdate msg <$> keyedHashInit key
{-|
Untagged variant of 'keyedHash'; takes a (possible 'undefined') key typed value from a 'KeyedHashAlgorithm' instance as parameter.

Example:

> keyedHash' (undefined :: HMAC SHA256) (fromString "secretkey") (fromString "secret message")
-}
keyedHash' :: KeyedHashAlgorithm k => k -> B.ByteString -> B.ByteString -> B.ByteString
keyedHash' k key msg = keyedHash key msg `witness` k
{-|
Helper to hash @key@ and lazy @message@ in one step

Example:

> untag (keyedHashLazy (fromString "secretkey") (fromString "secret message") :: Tagged (HMAC SHA256) B.ByteString)
-}
keyedHashLazy :: KeyedHashAlgorithm k => B.ByteString -> L.ByteString -> Tagged k B.ByteString
keyedHashLazy key msg = keyedHashFinalize <$> flip keyedHashUpdateLazy msg <$> keyedHashInit key
{-|
Untagged variant of 'keyedHashLazy'; takes a (possible 'undefined') key typed value from a 'KeyedHashAlgorithm' instance as parameter.

Example:

> keyedHashLazy' (undefined :: HMAC SHA256) (fromString "secretkey") (fromString "secret message")
-}
keyedHashLazy' :: KeyedHashAlgorithm k => k -> B.ByteString -> L.ByteString -> B.ByteString
keyedHashLazy' k key msg = keyedHashLazy key msg `witness` k

{-|
'HMAC' is a generic 'KeyedHashAlgorithm' instance to calculate the 'HMAC' based
on a 'HashAlgorithm'
-}
data HMAC a = HMAC !a !a

padZero :: Int -> B.ByteString -> B.ByteString
padZero len s = if len > B.length s then B.append s $ B.replicate (len - B.length s) 0 else s

instance HashAlgorithm a => KeyedHashAlgorithm (HMAC a) where
	implKeyedHashDigestSize = rt hashDigestSize where
		rt :: HashAlgorithm a => Tagged a x -> Tagged (HMAC a) x
		rt = retag
	implKeyedHashName = rt $ ("HMAC-" ++) <$> hashName where
		rt :: HashAlgorithm a => Tagged a x -> Tagged (HMAC a) x
		rt = retag
	implKeyedHashInit key = untag $ tagSelf hashInit >>= \i -> do
		blockSize <- hashBlockSize
		let key' = padZero blockSize $ if B.length key > blockSize then hash' i key else key
		let o_key = B.map (xor 0x5c) key'
		let i_key = B.map (xor 0x36) key'
		return $ HMAC (hashUpdate i o_key) (hashUpdate i i_key)
	implKeyedHashUpdate (HMAC o i) = HMAC o . hashUpdate i
	implKeyedHashUpdateLazy (HMAC o i) = HMAC o . hashUpdateLazy i
	implKeyedHashFinalize (HMAC o i) = hashFinalize $ hashUpdate o $ hashFinalize i

{-|
'hmacInit' is the default implementation for 'hashHMAC' and initializes a 'KeyedHash' to calculate
the HMAC for a message with the given @key@.

Example:

> let c = untag (hmacInit (fromString "secretkey") :: Tagged SHA256 KeyedHash) in keyedHashFinalize $ keyedHashUpdate c (fromString "secret message")
-}
hmacInit :: HashAlgorithm a => B.ByteString {- ^ @key@ argument -} -> Tagged a KeyedHash
hmacInit = rt . keyedHashInit where
	rt :: Tagged (HMAC a) x -> Tagged a x
	rt = retag

{-|
Untagged variant of 'hmacInit'; takes a (possible 'undefined') typed 'HashAlgorithm' context as parameter.

Example:

> keyedHashFinalize $ flip keyedHashUpdate (fromString "secret message") $ hmacInit' (undefined :: SHA256) (fromString "secretkey")
-}
hmacInit' :: HashAlgorithm a => a -> B.ByteString -> KeyedHash
hmacInit' a key = hmacInit key `witness` a

{-|
calculate HMAC with a 'HashAlgorithm' for a @key@ and @message@

Example:

> untag (hmac (fromString "secretkey") (fromString "secret message") :: Tagged SHA256 B.ByteString)
-}
hmac :: HashAlgorithm a => B.ByteString {- ^ @key@ argument -} -> B.ByteString {- ^ @message@ argument -} -> Tagged a B.ByteString
hmac key = rt . keyedHash key where
	rt :: Tagged (HMAC a) x -> Tagged a x
	rt = retag

{-|
Untagged variant of 'hmac'; takes a (possible 'undefined') typed 'HashAlgorithm' context as parameter.

Example:

> hmac' (undefined :: SHA256) (fromString "secretkey") (fromString "secret message")
-}
hmac' :: HashAlgorithm a => a -> B.ByteString -> B.ByteString -> B.ByteString
hmac' a key msg = hmac key msg `witness` a

{-|
calculate HMAC with a 'HashAlgorithm' for a @key@ and lazy @message@

Example:

> untag (hmacLazy (fromString "secretkey") (fromString "secret message") :: Tagged SHA256 B.ByteString)
-}
hmacLazy :: HashAlgorithm a => B.ByteString {- ^ @key@ argument -} -> L.ByteString {- ^ @message@ argument -} -> Tagged a B.ByteString
hmacLazy key = rt . keyedHashLazy key where
	rt :: Tagged (HMAC a) x -> Tagged a x
	rt = retag

{-|
Untagged variant of 'hmacLazy'; takes a (possible 'undefined') typed 'HashAlgorithm' context as parameter.

Example:

> hmacLazy' (undefined :: SHA256) (fromString "secretkey") (fromString "secret message")
-}
hmacLazy' :: HashAlgorithm a => a -> B.ByteString -> L.ByteString -> B.ByteString
hmacLazy' a key msg = hmacLazy key msg `witness` a