{-# LANGUAGE ScopedTypeVariables #-}

{-|
Description: Encryption of bytestrings using a type level nonce for determinism
License: BSD3

Given a value of a serializable type (like 'Int') we perform serialization and
compute a cryptographic hash of the associated namespace (carried as a phantom
type of kind 'Symbol').
The serialized payload is then encrypted using the symmetric cipher in CBC mode
using the hashed namespace as an initialization vector (IV).

The probability of detecting a namespace mismatch is thus \(1 - 2^{128-l}\)
where \(l\) is the length of the serialized payload.
-}
module Data.CryptoID.Poly
  ( CryptoID(..)
  , CryptoIDKey
  , genKey
  , encrypt
  , decrypt
  , CryptoIDError(..)
  , CryptoCipher, CryptoHash
  ) where

import Data.CryptoID

import Data.Binary
import Data.Binary.Put
import Data.Binary.Get

import Data.ByteString (ByteString)
import qualified Data.ByteString as ByteString
import qualified Data.ByteString.Char8 as ByteString.Char

import qualified Data.ByteString.Lazy as Lazy (ByteString)

import Data.List (sortOn)
import Data.Ord (Down(..))

import Data.ByteArray (ByteArrayAccess)
import qualified Data.ByteArray as ByteArray

import Data.Foldable (asum)
import Control.Monad.Except
import Control.Exception

import Data.Typeable
import GHC.TypeLits

import Crypto.Cipher.Types
import Crypto.Cipher.Blowfish (Blowfish)
import Crypto.Hash (hash, Digest)
import Crypto.Hash.Algorithms (SHAKE128)
import Crypto.Error

import Crypto.Random.Entropy


-- | The symmetric cipher 'BlockCipher' this module uses 
type CryptoCipher = Blowfish
-- | The cryptographic 'HashAlgorithm' this module uses
--
-- We expect the block size of 'CryptoCipher' to be exactly the size of the
-- 'Digest' generated by 'CryptoHash' (since a 'Digest' is used as an 'IV').
--
-- Violation of this expectation causes runtime errors.
type CryptoHash   = SHAKE128 64
  

-- | This newtype ensures only keys of the correct length can be created
--
-- Use 'genKey' to securely generate keys.
--
-- Use the 'Binary' instance to save and restore values of 'CryptoIDKey' across
-- executions.
newtype CryptoIDKey = CryptoIDKey { keyMaterial :: ByteString }
  deriving (Typeable, ByteArrayAccess)

-- | Does not actually show any key material
instance Show CryptoIDKey where
  show = show . typeOf

instance Binary CryptoIDKey where
  put = putByteString . keyMaterial
  get = CryptoIDKey <$> getKey (cipherKeySize cipher)
    where
      cipher :: CryptoCipher
      cipher = undefined

      -- Try key sizes from large to small ('Get' commits to the first branch
      -- that parses)
      getKey (KeySizeFixed n) = getByteString n
      getKey (KeySizeEnum ns) = asum [ getKey $ KeySizeFixed n | n <- sortOn Down ns ]
      getKey (KeySizeRange min max) = getKey $ KeySizeEnum [max .. min]


-- | Error cases that can be encountered during 'encrypt' and 'decrypt'
data CryptoIDError
  = AlgorithmError CryptoError
    -- ^ One of the underlying cryptographic algorithms
    --   ('CryptoHash' or 'CryptoCipher') failed.
  | NamespaceHashIsWrongLength ByteString
    -- ^ The length of the digest produced by 'CryptoHash' does
    --   not match the block size of 'CryptoCipher'.
    --
    -- The offending digest is included.
    --
    -- This error should not occur and is included primarily
    -- for sake of totality.
  | CiphertextConversionFailed
    -- ^ The produced 'ByteString' is the wrong length for conversion into a
    --   ciphertext.
  | DeserializationError (Lazy.ByteString, ByteOffset, String)
    -- ^ The plaintext obtained by decrypting a ciphertext with the given
    --   'CryptoIDKey' in the context of the @namespace@ could not be
    --   deserialized into a value of the expected @payload@-type.
    --
    -- This is expected behaviour if the @namespace@ or @payload@-type does not
    -- match the ones used during 'encrypt'ion or if the 'ciphertext' was
    -- tempered with.
  | InvalidNamespaceDetected
    -- ^ We have determined that, allthough deserializion succeded, the
    --   ciphertext was likely modified during transit or created using a
    --   different namespace.
  deriving (Show, Eq)

instance Exception CryptoIDError

-- | Securely generate a new key using system entropy
--
-- When 'CryptoCipher' accepts keys of varying lengths this function generates a
-- key of the largest accepted size.
genKey :: MonadIO m => m CryptoIDKey
genKey = CryptoIDKey <$> liftIO (getEntropy keySize)
  where
    keySize' = cipherKeySize (undefined :: CryptoCipher)

    keySize
      | KeySizeFixed n <- keySize' = n
      | KeySizeEnum ns <- keySize' = maximum ns
      | KeySizeRange _ max <- keySize' = max
  
  
-- | @pad err size src@ appends null bytes to @src@ until it has length that is
--   a multiple of @size@.
pad :: ByteArrayAccess a => Int -> a -> ByteString
pad n (ByteArray.unpack -> src) = ByteString.pack $ src ++ replicate (l `mod` n) 0
  where
    l = length src

-- | Use 'CryptoHash' to generate a 'Digest' of the Symbol passed as proxy type
namespace' :: forall proxy namespace m.
              ( KnownSymbol namespace, MonadError CryptoIDError m
              ) => proxy namespace -> m (IV CryptoCipher)
namespace' p = case makeIV namespaceHash of
                 Nothing -> throwError . NamespaceHashIsWrongLength $ ByteArray.convert namespaceHash
                 Just iv -> return iv
  where
    namespaceHash :: Digest CryptoHash
    namespaceHash = hash . ByteString.Char.pack $ symbolVal p

-- | Wrap failure of one of the cryptographic algorithms as a 'CryptoIDError'
cryptoFailable :: MonadError CryptoIDError m => CryptoFailable a -> m a
cryptoFailable = either (throwError . AlgorithmError) return . eitherCryptoError

-- | Encrypt an arbitrary serializable value
encrypt :: forall m namespace.
           ( KnownSymbol namespace
           , MonadError CryptoIDError m
           ) => CryptoIDKey -> ByteString -> m (CryptoID namespace ByteString)
encrypt (keyMaterial -> key) plaintext = do
  cipher <- cryptoFailable (cipherInit key :: CryptoFailable CryptoCipher)
  namespace <- namespace' (Proxy :: Proxy namespace)
  return . CryptoID . cbcEncrypt cipher namespace $ pad (blockSize cipher) plaintext


-- | Decrypt an arbitrary serializable value
decrypt :: forall m namespace.
           ( KnownSymbol namespace
           , MonadError CryptoIDError m
           ) => CryptoIDKey -> CryptoID namespace ByteString -> m ByteString
decrypt (keyMaterial -> key) CryptoID{..} = do
  cipher <- cryptoFailable (cipherInit key :: CryptoFailable CryptoCipher)
  namespace <- namespace' (Proxy :: Proxy namespace)
  return $ cbcDecrypt cipher namespace ciphertext