{-# LANGUAGE ScopedTypeVariables #-}

{-|
Description: Reversably generate UUIDs from arbitrary serializable types in a secure fashion
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 a symmetric cipher in CBC
mode using the hashed namespace as an initialization vector (IV).

Since the serialized payload is padded to exactly the length of a single cipher
block we can detect namespace mismatches by checking that all bytes expected to
have been inserted during padding are nil.
The probability of detecting a namespace mismatch is thus \(1 - 2^{128-l}\)
where \(l\) is the length of the serialized payload.
-}
module Data.UUID.Cryptographic
  ( CryptoID(..)
  , CryptoIDKey
  , genKey
  , CryptoIDError(..)
  , encrypt
  , decrypt
  , CryptoCipher, CryptoHash
  ) where

import Data.UUID (UUID, toByteString, fromByteString)
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 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.AES (AES256)
import Crypto.Hash (hash, Digest)
import Crypto.Hash.Algorithms (SHAKE128)
import Crypto.Error

import Crypto.Random.Entropy

import GHC.Generics (Generic)
import Data.Data (Data)
import Foreign.Storable (Storable)


-- | The symmetric cipher 'BlockCipher' this module uses 
type CryptoCipher = AES256
-- | The cryptographic 'HashAlgorithm' this module uses
--
-- We expect the block size of 'CryptoCipher' to be exactly the size of the
-- 'Digest' generated by 'CryptoHash'.
--
-- Violation of this expectation causes runtime errors.
type CryptoHash   = SHAKE128 128
  

-- | 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]


-- | A thin wrapper around 'UUID' to carry the type information from which we
--   infer what payload we expect the 'UUID' to contain.
newtype CryptoID (namespace :: Symbol) = CryptoID { ciphertext :: UUID }
  deriving (Eq, Data, Ord, Read, Show, Binary, Typeable, Storable, Generic)

-- | 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.
                   | PlaintextTooLong ByteString
                     -- ^ The serialized representation of the payload exceeds
                     --   the block size of 'CryptoCipher'.
                     --
                     -- The offending serialization is included.
                   | UUIDConversionFailed
                     -- ^ The length of the produced ciphertext (which is
                     --   expected to be exactly one block size of
                     --   'CryptoCipher') does not match the size expected for a
                     --   'UUID' (128 bits).
                     --
                     -- This error should not occur and is included primarily
                     -- for sake of totality.
                   | DeserializationError (Lazy.ByteString, ByteOffset, String)
                     -- ^ The plaintext obtained by decrypting a 'UUID' 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 'UUID' 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 @size@.
--
-- If @src@ is already longer than @size@ @err@ is thrown instead.
pad :: (MonadError e m, ByteArrayAccess a) => (ByteString -> e) -> Int -> a -> m ByteString
pad err n (ByteArray.unpack -> src)
  | l > n     = throwError . err $ ByteString.pack src
  | otherwise = return . ByteString.pack $ src ++ replicate (n - l) 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 = hash . ByteString.Char.pack $ symbolVal p :: Digest CryptoHash

-- | 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
--
-- We only expect to fail if the given value is not serialized in such a fashion
-- that it fits within one 'CryptoCipher'-block.
--
-- Larger values could likely not be contained wholly within 128 bits (the size
-- of an 'UUID') in any case.
encrypt :: forall a m namespace.
           ( KnownSymbol namespace
           , Binary a
           , MonadError CryptoIDError m
           ) => CryptoIDKey -> a -> m (CryptoID namespace)
encrypt (keyMaterial -> key) val = do
  cipher <- cryptoFailable (cipherInit key :: CryptoFailable CryptoCipher)
  
  namespace <- namespace' (Proxy :: Proxy namespace)
  plaintext <- pad PlaintextTooLong (blockSize cipher) . Lazy.ByteString.toStrict $ encode val
  
  CryptoID <$> uuidConversion (cbcEncrypt cipher namespace plaintext)
    where
      uuidConversion = maybe (throwError UUIDConversionFailed) return . fromByteString . Lazy.ByteString.fromStrict


-- | Decrypt an arbitrary serializable value
--
-- Since no integrity guarantees can be made (we do not sign the values we
-- 'encrypt') it is likely that deserialization will fail emitting
-- 'DeserializationError' or 'InvalidNamespaceDetected'.
decrypt :: forall a m namespace.
           ( KnownSymbol namespace
           , Binary a
           , MonadError CryptoIDError m
           ) => CryptoIDKey -> CryptoID namespace -> m a
decrypt (keyMaterial -> key) CryptoID{..} = do
  cipher <- cryptoFailable (cipherInit key :: CryptoFailable CryptoCipher)
  
  namespace <- namespace' (Proxy :: Proxy namespace)
  let ciphertext' = Lazy.ByteString.toStrict $ toByteString ciphertext
      plaintext = Lazy.ByteString.fromStrict (cbcDecrypt cipher namespace ciphertext')

  case decodeOrFail plaintext of
    Left err -> throwError $ DeserializationError err
    Right (rem, _, res)
      | Lazy.ByteString.all (== 0) rem -> return res
      | otherwise -> throwError InvalidNamespaceDetected