-- |

-- Module      : Crypto.Data.Padding

-- License     : BSD-style

-- Maintainer  : Vincent Hanquez <vincent@snarc.org>

-- Stability   : experimental

-- Portability : unknown

--

-- Various cryptographic padding commonly used for block ciphers

-- or assymetric systems.

--

module Crypto.Data.Padding

    ( Format(..)

    , pad

    , unpad

    ) where

import           Data.ByteArray (ByteArray, Bytes)

import qualified Data.ByteArray as B

-- | Format of padding

data Format =

      PKCS5     -- ^ PKCS5: PKCS7 with hardcoded size of 8

    | PKCS7 Int -- ^ PKCS7 with padding size between 1 and 255

    | ZERO Int  -- ^ zero padding with block size

    deriving (Show, Eq)

-- | Apply some pad to a bytearray

pad :: ByteArray byteArray => Format -> byteArray -> byteArray

pad  PKCS5     bin = pad (PKCS7 8) bin

pad (PKCS7 sz) bin = bin `B.append` paddingString

  where

    paddingString = B.replicate paddingByte (fromIntegral paddingByte)

    paddingByte   = sz - (B.length bin `mod` sz)

pad (ZERO sz)  bin = bin `B.append` paddingString

  where

    paddingString = B.replicate paddingSz 0

    paddingSz

      | len == 0   =  sz

      | m == 0     =  0

      | otherwise  =  sz - m

    m = len `mod` sz

    len = B.length bin

-- | Try to remove some padding from a bytearray.

unpad :: ByteArray byteArray => Format -> byteArray -> Maybe byteArray

unpad  PKCS5     bin = unpad (PKCS7 8) bin

unpad (PKCS7 sz) bin

    | len == 0                           = Nothing

    | (len `mod` sz) /= 0                = Nothing

    | paddingSz < 1 || paddingSz > len   = Nothing

    | paddingWitness `B.constEq` padding = Just content

    | otherwise                          = Nothing

  where

    len         = B.length bin

    paddingByte = B.index bin (len - 1)

    paddingSz   = fromIntegral paddingByte

    (content, padding) = B.splitAt (len - paddingSz) bin

    paddingWitness     = B.replicate paddingSz paddingByte :: Bytes

unpad (ZERO sz)  bin

    | len == 0                           = Nothing

    | (len `mod` sz) /= 0                = Nothing

    | B.index bin (len - 1) /= 0         = Just bin

    | otherwise                          = Nothing

  where

    len         = B.length bin