{-# LANGUAGE DeriveDataTypeable #-}

-- |
-- Module      : Crypto.PubKey.Rabin.RW
-- License     : BSD-style
-- Maintainer  : Carlos Rodriguez-Vega <crodveg@yahoo.es>
-- Stability   : experimental
-- Portability : unknown
--
-- Rabin-Williams cryptosystem for public-key encryption and digital signature.
-- See pages 323 - 324 in "Computational Number Theory and Modern Cryptography" by Song Y. Yan.
-- Also inspired by https://github.com/vanilala/vncrypt/blob/master/vncrypt/vnrw_gmp.c.
module Crypto.PubKey.Rabin.RW (
    PublicKey (..),
    PrivateKey (..),
    generate,
    encrypt,
    encryptWithSeed,
    decrypt,
    sign,
    verify,
) where

import Data.ByteString
import Data.Data

import Crypto.Hash
import Crypto.Number.Basic (numBytes)
import Crypto.Number.ModArithmetic (expSafe, jacobi)
import Crypto.Number.Serialize (i2osp, i2ospOf_, os2ip)
import Crypto.PubKey.Rabin.OAEP
import Crypto.PubKey.Rabin.Types
import Crypto.Random.Types

-- | Represent a Rabin-Williams public key.
data PublicKey = PublicKey
    { public_size :: Int
    -- ^ size of key in bytes
    , public_n :: Integer
    -- ^ public p*q
    }
    deriving (Show, Read, Eq, Data)

-- | Represent a Rabin-Williams private key.
data PrivateKey = PrivateKey
    { private_pub :: PublicKey
    , private_p :: Integer
    -- ^ p prime number
    , private_q :: Integer
    -- ^ q prime number
    , private_d :: Integer
    }
    deriving (Show, Read, Eq, Data)

-- | Generate a pair of (private, public) key of size in bytes.
-- Prime p is congruent 3 mod 8 and prime q is congruent 7 mod 8.
generate
    :: MonadRandom m
    => Int
    -> m (PublicKey, PrivateKey)
generate size = do
    (p, q) <- generatePrimes size (\p -> p `mod` 8 == 3) (\q -> q `mod` 8 == 7)
    return (generateKeys p q)
  where
    generateKeys p q =
        let n = p * q
            d = ((p - 1) * (q - 1) `div` 4 + 1) `div` 2
            publicKey =
                PublicKey
                    { public_size = size
                    , public_n = n
                    }
            privateKey =
                PrivateKey
                    { private_pub = publicKey
                    , private_p = p
                    , private_q = q
                    , private_d = d
                    }
         in (publicKey, privateKey)

-- | Encrypt plaintext using public key an a predefined OAEP seed.
--
-- See algorithm 8.11 in "Handbook of Applied Cryptography" by Alfred J. Menezes et al.
encryptWithSeed
    :: HashAlgorithm hash
    => ByteString
    -- ^ Seed
    -> OAEPParams hash ByteString ByteString
    -- ^ OAEP padding
    -> PublicKey
    -- ^ public key
    -> ByteString
    -- ^ plaintext
    -> Either Error ByteString
encryptWithSeed seed oaep pk m =
    let n = public_n pk
        k = numBytes n
     in do
            m' <- pad seed oaep k m
            m'' <- ep1 n $ os2ip m'
            return $ i2osp $ ep2 n m''

-- | Encrypt plaintext using public key.
encrypt
    :: (HashAlgorithm hash, MonadRandom m)
    => OAEPParams hash ByteString ByteString
    -- ^ OAEP padding parameters
    -> PublicKey
    -- ^ public key
    -> ByteString
    -- ^ plaintext
    -> m (Either Error ByteString)
encrypt oaep pk m = do
    seed <- getRandomBytes hashLen
    return $ encryptWithSeed seed oaep pk m
  where
    hashLen = hashDigestSize (oaepHash oaep)

-- | Decrypt ciphertext using private key.
decrypt
    :: HashAlgorithm hash
    => OAEPParams hash ByteString ByteString
    -- ^ OAEP padding parameters
    -> PrivateKey
    -- ^ private key
    -> ByteString
    -- ^ ciphertext
    -> Maybe ByteString
decrypt oaep pk c =
    let d = private_d pk
        n = public_n $ private_pub pk
        k = numBytes n
        c' = i2ospOf_ k $ dp2 n $ dp1 d n $ os2ip c
     in case unpad oaep k c' of
            Left _ -> Nothing
            Right p -> Just p

-- | Sign message using hash algorithm and private key.
sign
    :: HashAlgorithm hash
    => PrivateKey
    -- ^ private key
    -> hash
    -- ^ hash function
    -> ByteString
    -- ^ message to sign
    -> Either Error Integer
sign pk hashAlg m =
    let d = private_d pk
        n = public_n $ private_pub pk
     in do
            m' <- ep1 n $ os2ip $ hashWith hashAlg m
            return $ dp1 d n m'

-- | Verify signature using hash algorithm and public key.
verify
    :: HashAlgorithm hash
    => PublicKey
    -- ^ public key
    -> hash
    -- ^ hash function
    -> ByteString
    -- ^ message
    -> Integer
    -- ^ signature
    -> Bool
verify pk hashAlg m s =
    let n = public_n pk
        h = os2ip $ hashWith hashAlg m
        h' = dp2 n $ ep2 n s
     in h' == h

-- | Encryption primitive 1
ep1 :: Integer -> Integer -> Either Error Integer
ep1 n m =
    let m' = 2 * m + 1
        m'' = 2 * m'
        m''' = 2 * m''
     in case jacobi m' n of
            Just (-1) | m'' < n -> Right m''
            Just 1 | m''' < n -> Right m'''
            _ -> Left InvalidParameters

-- | Encryption primitive 2
ep2 :: Integer -> Integer -> Integer
ep2 n m = expSafe m 2 n

-- | Decryption primitive 1
dp1 :: Integer -> Integer -> Integer -> Integer
dp1 d n c = expSafe c d n

-- | Decryption primitive 2
dp2 :: Integer -> Integer -> Integer
dp2 n c =
    let c' = c `div` 2
        c'' = (n - c) `div` 2
     in case c `mod` 4 of
            0 -> ((c' `div` 2 - 1) `div` 2)
            1 -> ((c'' `div` 2 - 1) `div` 2)
            2 -> ((c' - 1) `div` 2)
            _ -> ((c'' - 1) `div` 2)
