Simple AES-GCM encrypt-decrypt

{-# LANGUAGE PackageImports #-}
module AES (encryptFile, decrypt) where

import Data.ByteString (ByteString)
import qualified Data.ByteString as B
import Data.Maybe
import qualified "cryptonite" Crypto.Cipher.AES as C
import qualified "cryptonite" Crypto.Cipher.Types as C
import qualified "cryptonite" Crypto.Error as C
import qualified Crypto.Nonce as N
import System.IO (withFile, IOMode(..))
import qualified Data.ByteArray as BA
import qualified Data.ByteString.Lazy as LB
import qualified Control.Exception as E
import Data.Typeable (Typeable)

data AESError = CryptoError String | AuthTagFailed deriving (Show, Typeable)
instance E.Exception AESError

bufSize :: Int
bufSize = 65536

authTagLength :: Int
authTagLength = 16

withCipher :: (ByteString, ByteString) -> (C.AEAD C.AES128 -> a) -> a
withCipher (ek,iv) f = case C.cipherInit ek of
    C.CryptoPassed aes -> case C.aeadInit C.AEAD_GCM aes iv of
        C.CryptoPassed aead -> f aead
        C.CryptoFailed e -> E.throw $ CryptoError $ show e

mkAuthTag :: C.Cipher cipher => C.AEAD cipher -> ByteString
mkAuthTag ctx = BA.convert $ C.unAuthTag $ C.aeadFinalize ctx authTagLength

encryptFile :: FilePath -> FilePath -> IO (ByteString, ByteString, ByteString)
encryptFile srcPath dstPath = withFile srcPath ReadMode $ \src ->
    withFile dstPath WriteMode $ \dst -> do
        g <- N.new
        ek <- N.nonce128 g
        iv <- N.nonce128 g
        withCipher (ek,iv) $ \aead -> let
                loop ctx = do
                    buf <- B.hGet src bufSize
                    if B.null buf
                        then return (ek, iv, mkAuthTag ctx)
                        else do
                            let (c, ctx') = C.aeadEncrypt ctx buf
                            B.hPut dst c
                            loop ctx'
            in loop aead

decrypt :: (ByteString, ByteString, ByteString) -> LB.ByteString -> LB.ByteString
decrypt (ek,iv,at) c = withCipher (ek,iv) $ \aead -> let
        loop ctx (c:cs) = let (p, ctx') = C.aeadDecrypt ctx c in p:loop ctx' cs
        loop ctx []     = if at == mkAuthTag ctx
                then []
                else E.throw AuthTagFailed
    in LB.fromChunks $ loop aead $ LB.toChunks c