    -- |
    -- Module      : Crypto.PubKey.RSA.PKCS15
    -- License     : BSD-style
    -- Maintainer  : Vincent Hanquez <vincent@snarc.org>
    -- Stability   : experimental
    -- Portability : Good
    --
    module Crypto.PubKey.RSA.PKCS15
        (
        -- * padding and unpadding
          pad
        , padSignature
        , unpad
        -- * private key operations
        , decrypt
        , decryptSafer
        , sign
        , signSafer
        -- * public key operations
        , encrypt
        , verify
        -- * hash ASN1 description
        , HashAlgorithmASN1
        ) where

    import           Crypto.Random.Types
    import           Crypto.PubKey.Internal (and')
    import           Crypto.PubKey.RSA.Types
    import           Crypto.PubKey.RSA.Prim
    import           Crypto.PubKey.RSA (generateBlinder)
    import           Crypto.Hash

    import           Data.ByteString (ByteString)
    import           Data.Word

    import           Crypto.Internal.ByteArray (ByteArray, Bytes)
    import qualified Crypto.Internal.ByteArray as B

    -- | A specialized class for hash algorithm that can product
    -- a ASN1 wrapped description the algorithm plus the content
    -- of the digest.
    class HashAlgorithm hashAlg => HashAlgorithmASN1 hashAlg where
        -- | Convert a Digest into an ASN1 wrapped descriptive ByteArray
        hashDigestASN1 :: ByteArray out => Digest hashAlg -> out

    -- http://uk.emc.com/emc-plus/rsa-labs/pkcs/files/h11300-wp-pkcs-1v2-2-rsa-cryptography-standard.pdf
    -- EMSA-PKCS1-v1_5
    instance HashAlgorithmASN1 MD2 where
        hashDigestASN1 = addDigestPrefix [0x30,0x20,0x30,0x0c,0x06,0x08,0x2a,0x86,0x48,0x86,0xf7,0x0d,0x02,0x02,0x05,0x00,0x04,0x10]
    instance HashAlgorithmASN1 MD5 where
        hashDigestASN1 = addDigestPrefix [0x30,0x20,0x30,0x0c,0x06,0x08,0x2a,0x86,0x48,0x86,0xf7,0x0d,0x02,0x05,0x05,0x00,0x04,0x10]
    instance HashAlgorithmASN1 SHA1 where
        hashDigestASN1 = addDigestPrefix [0x30,0x21,0x30,0x09,0x06,0x05,0x2b,0x0e,0x03,0x02,0x1a,0x05,0x00,0x04,0x14]
    instance HashAlgorithmASN1 SHA224 where
        hashDigestASN1 = addDigestPrefix [0x30,0x2d,0x30,0x0d,0x06,0x09,0x60,0x86,0x48,0x01,0x65,0x03,0x04,0x02,0x04,0x05,0x00,0x04,0x1c]
    instance HashAlgorithmASN1 SHA256 where
        hashDigestASN1 = addDigestPrefix [0x30,0x31,0x30,0x0d,0x06,0x09,0x60,0x86,0x48,0x01,0x65,0x03,0x04,0x02,0x01,0x05,0x00,0x04,0x20]
    instance HashAlgorithmASN1 SHA384 where
        hashDigestASN1 = addDigestPrefix [0x30,0x41,0x30,0x0d,0x06,0x09,0x60,0x86,0x48,0x01,0x65,0x03,0x04,0x02,0x02,0x05,0x00,0x04,0x30]
    instance HashAlgorithmASN1 SHA512 where
        hashDigestASN1 = addDigestPrefix [0x30,0x51,0x30,0x0d,0x06,0x09,0x60,0x86,0x48,0x01,0x65,0x03,0x04,0x02,0x03,0x05,0x00,0x04,0x40]
    instance HashAlgorithmASN1 SHA512t_224 where
        hashDigestASN1 = addDigestPrefix [0x30,0x2d,0x30,0x0d,0x06,0x09,0x60,0x86,0x48,0x01,0x65,0x03,0x04,0x02,0x05,0x05,0x00,0x04,0x1c]
    instance HashAlgorithmASN1 SHA512t_256 where
        hashDigestASN1 = addDigestPrefix [0x30,0x31,0x30,0x0d,0x06,0x09,0x60,0x86,0x48,0x01,0x65,0x03,0x04,0x02,0x06,0x05,0x00,0x04,0x20]
    instance HashAlgorithmASN1 RIPEMD160 where
        hashDigestASN1 = addDigestPrefix [0x30,0x21,0x30,0x09,0x06,0x05,0x2b,0x24,0x03,0x02,0x01,0x05,0x00,0x04,0x14]

    --
    -- ** Hack **
    --
    -- this happens to not need a real ASN1 encoder, because
    -- thanks to the digest being a specific size AND
    -- that the digest data is the last bytes in the encoding,
    -- this allows to just prepend the right prefix to the
    -- computed digest, to make it in the expected and valid shape.
    --
    -- Otherwise the expected structure is in the following form:
    --
    --   Start Sequence
    --     ,Start Sequence
    --       ,OID oid
    --       ,Null
    --     ,End Sequence
    --     ,OctetString digest
    --   ,End Sequence
    addDigestPrefix :: ByteArray out => [Word8] -> Digest hashAlg -> out
    addDigestPrefix prefix digest =
        B.pack prefix `B.append` B.convert digest

    -- | This produce a standard PKCS1.5 padding for encryption
    pad :: (MonadRandom m, ByteArray message) => Int -> message -> m (Either Error message)
    pad len m
        | B.length m > len - 11 = return (Left MessageTooLong)
        | otherwise             = do
            padding <- getNonNullRandom (len - B.length m - 3)
            return $ Right $ B.concat [ B.pack [0,2], padding, B.pack [0], m ]

      where
        -- get random non-null bytes
        getNonNullRandom :: (ByteArray bytearray, MonadRandom m) => Int -> m bytearray
        getNonNullRandom n = do
            bs0 <- getRandomBytes n
            let bytes = B.pack $ filter (/= 0) $ B.unpack (bs0 :: Bytes)
                left  = n - B.length bytes
            if left == 0
                then return bytes
                else do bend <- getNonNullRandom left
                        return (bytes `B.append` bend)

    -- | Produce a standard PKCS1.5 padding for signature
    padSignature :: ByteArray signature => Int -> signature -> Either Error signature
    padSignature klen signature
        | klen < siglen+1 = Left SignatureTooLong
        | otherwise       = Right (B.pack padding `B.append` signature)
      where
            siglen    = B.length signature
            padding   = 0 : 1 : (replicate (klen - siglen - 3) 0xff ++ [0])

    -- | Try to remove a standard PKCS1.5 encryption padding.
    unpad :: ByteArray bytearray => bytearray -> Either Error bytearray
    unpad packed
        | paddingSuccess = Right m
        | otherwise      = Left MessageNotRecognized
      where
            (zt, ps0m)   = B.splitAt 2 packed
            (ps, zm)     = B.span (/= 0) ps0m
            (z, m)       = B.splitAt 1 zm
            paddingSuccess = and' [ zt `B.constEq` (B.pack [0,2] :: Bytes)
                                  , z == B.zero 1
                                  , B.length ps >= 8
                                  ]

    -- | decrypt message using the private key.
    --
    -- When the decryption is not in a context where an attacker could gain
    -- information from the timing of the operation, the blinder can be set to None.
    --
    -- If unsure always set a blinder or use decryptSafer
    decrypt :: Maybe Blinder -- ^ optional blinder
            -> PrivateKey    -- ^ RSA private key
            -> ByteString    -- ^ cipher text
            -> Either Error ByteString
    decrypt blinder pk c
        | B.length c /= (private_size pk) = Left MessageSizeIncorrect
        | otherwise                       = unpad $ dp blinder pk c

    -- | decrypt message using the private key and by automatically generating a blinder.
    decryptSafer :: MonadRandom m
                 => PrivateKey -- ^ RSA private key
                 -> ByteString -- ^ cipher text
                 -> m (Either Error ByteString)
    decryptSafer pk b = do
        blinder <- generateBlinder (private_n pk)
        return (decrypt (Just blinder) pk b)

    -- | encrypt a bytestring using the public key.
    --
    -- the message needs to be smaller than the key size - 11
    encrypt :: MonadRandom m => PublicKey -> ByteString -> m (Either Error ByteString)
    encrypt pk m = do
        r <- pad (public_size pk) m
        case r of
            Left err -> return $ Left err
            Right em -> return $ Right (ep pk em)

    -- | sign message using private key, a hash and its ASN1 description
    --
    -- When the signature is not in a context where an attacker could gain
    -- information from the timing of the operation, the blinder can be set to None.
    --
    -- If unsure always set a blinder or use signSafer
    sign :: HashAlgorithmASN1 hashAlg
         => Maybe Blinder -- ^ optional blinder
         -> Maybe hashAlg -- ^ hash algorithm
         -> PrivateKey    -- ^ private key
         -> ByteString    -- ^ message to sign
         -> Either Error ByteString
    sign blinder hashDescr pk m = dp blinder pk `fmap` makeSignature hashDescr (private_size pk) m

    -- | sign message using the private key and by automatically generating a blinder.
    signSafer :: (HashAlgorithmASN1 hashAlg, MonadRandom m)
              => Maybe hashAlg -- ^ Hash algorithm
              -> PrivateKey    -- ^ private key
              -> ByteString    -- ^ message to sign
              -> m (Either Error ByteString)
    signSafer hashAlg pk m = do
        blinder <- generateBlinder (private_n pk)
        return (sign (Just blinder) hashAlg pk m)

    -- | verify message with the signed message
    verify :: HashAlgorithmASN1 hashAlg
           => Maybe hashAlg
           -> PublicKey
           -> ByteString
           -> ByteString
           -> Bool
    verify hashAlg pk m sm =
        case makeSignature hashAlg (public_size pk) m of
            Left _  -> False
            Right s -> s == (ep pk sm)

    -- | make signature digest, used in 'sign' and 'verify'
    makeSignature :: HashAlgorithmASN1 hashAlg
                  => Maybe hashAlg -- ^ optional hashing algorithm
                  -> Int
                  -> ByteString
                  -> Either Error ByteString
    makeSignature Nothing        klen m = padSignature klen m
    makeSignature (Just hashAlg) klen m = padSignature klen (hashDigestASN1 $ hashWith hashAlg m)