-- CFB.hs: OpenPGP (RFC4880) CFB mode
-- Copyright © 2013 Daniel Kahn Gillmor and Clint Adams
-- This software is released under the terms of the Expat license.
-- (See the LICENSE file).

module Codec.Encryption.OpenPGP.CFB (
   decrypt
 , decryptNoNonce
 , decryptOpenPGPCfb
 , encryptNoNonce
) where

import Codec.Encryption.OpenPGP.BlockCipher (BCipher(..), bcBlockSize)
import Codec.Encryption.OpenPGP.Types
import Control.Applicative ((<$>), (<*>))
import Control.Error.Util (note)
import Crypto.Cipher.Types (makeKey, nullIV, BlockCipher(..), Cipher(..))
import qualified Crypto.Cipher as CC
import qualified Crypto.Nettle.Ciphers as CNC
import qualified Data.ByteString as B
import Data.SecureMem (ToSecureMem)

decryptOpenPGPCfb :: SymmetricAlgorithm -> B.ByteString -> B.ByteString -> Either String B.ByteString
decryptOpenPGPCfb Plaintext ciphertext _ = return ciphertext
decryptOpenPGPCfb sa ciphertext keydata = do
    bc <- mkBCipher sa keydata
    let nonce = decrypt1 ciphertext bc
    cleartext <- decrypt2 ciphertext bc
    if nonceCheck bc nonce then return cleartext else fail "Session key quickcheck failed"
    where
        decrypt1 :: B.ByteString -> BCipher -> B.ByteString
        decrypt1 ct (BCipher cipher) = cdecrypt sa cipher nullIV (B.take (blockSize cipher + 2) ct)
        decrypt2 :: B.ByteString -> BCipher -> Either String B.ByteString
        decrypt2 ct (BCipher cipher) = note "unexpected CFB-resync failure" (CC.makeIV (B.take (blockSize cipher) (B.drop 2 ct))) >>= \i -> return (cdecrypt sa cipher i (B.drop (blockSize cipher + 2) ct))

decrypt :: SymmetricAlgorithm -> B.ByteString -> B.ByteString -> Either String B.ByteString
decrypt Plaintext ciphertext _ = return ciphertext
decrypt sa ciphertext keydata = do
    bc <- mkBCipher sa keydata
    let (nonce, cleartext) = B.splitAt (bcBlockSize bc + 2) (decrypt' ciphertext bc)
    if nonceCheck bc nonce then return cleartext else fail "Session key quickcheck failed"
    where
        decrypt' :: B.ByteString -> BCipher -> B.ByteString
        decrypt' ct (BCipher cipher) = cdecrypt sa cipher nullIV ct

decryptNoNonce :: SymmetricAlgorithm -> IV -> B.ByteString -> B.ByteString -> Either String B.ByteString
decryptNoNonce Plaintext _ ciphertext _ = return ciphertext
decryptNoNonce sa iv ciphertext keydata = do
    bc <- mkBCipher sa keydata
    decrypt' ciphertext bc
    where
        decrypt' :: B.ByteString -> BCipher -> Either String B.ByteString
        decrypt' ct (BCipher cipher) = note "Bad IV" (CC.makeIV iv) >>= \i -> return (cdecrypt sa cipher i ct)

cdecrypt :: BlockCipher cipher => SymmetricAlgorithm -> cipher -> CC.IV cipher -> B.ByteString -> B.ByteString
cdecrypt sa
    | sa `elem` [CAST5, Twofish] = paddedCfbDecrypt
    | otherwise = cfbDecrypt

nonceCheck :: BCipher -> B.ByteString -> Bool
nonceCheck bc = (==) <$> B.take 2 . B.drop (bcBlockSize bc - 2) <*> B.drop (bcBlockSize bc)

paddedCfbDecrypt :: BlockCipher cipher => cipher -> CC.IV cipher -> B.ByteString -> B.ByteString
paddedCfbDecrypt cipher iv ciphertext = B.take (B.length ciphertext) (cfbDecrypt cipher iv padded)
    where
        padded = ciphertext `B.append` B.pack (replicate (blockSize cipher - (B.length ciphertext `mod` blockSize cipher)) 0)

mkBCipher :: ToSecureMem b => SymmetricAlgorithm -> b -> Either String BCipher
mkBCipher Plaintext = const (fail "this shouldn't have happened") -- FIXME: orphan instance?
mkBCipher IDEA = const (fail "IDEA not yet implemented") -- FIXME: IDEA
mkBCipher ReservedSAFER = const (fail "SAFER not implemented") -- FIXME: or not?
mkBCipher ReservedDES = const (fail "DES not implemented") -- FIXME: or not?
mkBCipher (OtherSA _) = const (fail "Unknown, unimplemented symmetric algorithm")
mkBCipher CAST5 = return . BCipher . (ciph :: ToSecureMem b => b -> CNC.CAST128)
mkBCipher Twofish = return . BCipher . (ciph :: ToSecureMem b => b -> CNC.TWOFISH)
mkBCipher TripleDES = return . BCipher . (ciph :: ToSecureMem b => b -> CC.DES_EDE3)
mkBCipher Blowfish = return . BCipher . (ciph :: ToSecureMem b => b -> CC.Blowfish128)
mkBCipher AES128 = return . BCipher . (ciph :: ToSecureMem b => b -> CC.AES128)
mkBCipher AES192 = return . BCipher . (ciph :: ToSecureMem b => b -> CC.AES192)
mkBCipher AES256 = return . BCipher . (ciph :: ToSecureMem b => b -> CC.AES256)

ciph :: (CC.BlockCipher cipher, ToSecureMem b) => b -> cipher -- FIXME: return an Either String cipher ?
ciph keydata = cipherInit ekey
    where
        ekey = case makeKey keydata of
          Left _ -> error "bad cipher parameters"
          Right key -> key

encryptNoNonce :: SymmetricAlgorithm -> S2K -> IV -> B.ByteString -> B.ByteString -> Either String B.ByteString
encryptNoNonce Plaintext _ _ payload keydata = return payload
encryptNoNonce sa s2k iv payload keydata = do
    bc <- mkBCipher sa keydata
    encrypt' payload bc
    where
        encrypt' :: B.ByteString -> BCipher -> Either String B.ByteString
        encrypt' ct (BCipher cipher) = note "Bad IV" (CC.makeIV iv) >>= \i -> return (cencrypt sa cipher i ct)

cencrypt :: BlockCipher cipher => SymmetricAlgorithm -> cipher -> CC.IV cipher -> B.ByteString -> B.ByteString
cencrypt sa
    | sa `elem` [CAST5, Twofish] = error "padding for nettle-encryption not implemented yet"
    | otherwise = cfbEncrypt