{-# LANGUAGE DeriveDataTypeable #-}

-- |
-- Module      : Crypto.PubKey.Rabin.Basic
-- License     : BSD-style
-- Maintainer  : Carlos Rodriguez-Vega <crodveg@yahoo.es>
-- Stability   : experimental
-- Portability : unknown
--
-- Rabin cryptosystem for public-key cryptography and digital signature.
module Crypto.PubKey.Rabin.Basic (
    PublicKey (..),
    PrivateKey (..),
    Signature (..),
    generate,
    encrypt,
    encryptWithSeed,
    decrypt,
    sign,
    signWith,
    verify,
) where

import Data.ByteString (ByteString)
import qualified Data.ByteString as B
import Data.Data
import Data.Either (rights)

import Crypto.Hash
import Crypto.Number.Basic (gcde, 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 (MonadRandom, getRandomBytes)

-- | Represent a Rabin 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 private key.
data PrivateKey = PrivateKey
    { private_pub :: PublicKey
    , private_p :: Integer
    -- ^ p prime number
    , private_q :: Integer
    -- ^ q prime number
    , private_a :: Integer
    , private_b :: Integer
    }
    deriving (Show, Read, Eq, Data)

-- | Rabin Signature.
data Signature = Signature (Integer, Integer) deriving (Show, Read, Eq, Data)

-- | Generate a pair of (private, public) key of size in bytes.
-- Primes p and q are both congruent 3 mod 4.
--
-- See algorithm 8.11 in "Handbook of Applied Cryptography" by Alfred J. Menezes et al.
generate
    :: MonadRandom m
    => Int
    -> m (PublicKey, PrivateKey)
generate size = do
    (p, q) <- generatePrimes size (\p -> p `mod` 4 == 3) (\q -> q `mod` 4 == 3)
    return $ generateKeys p q
  where
    generateKeys p q =
        let n = p * q
            (a, b, _) = gcde p q
            publicKey =
                PublicKey
                    { public_size = size
                    , public_n = n
                    }
            privateKey =
                PrivateKey
                    { private_pub = publicKey
                    , private_p = p
                    , private_q = q
                    , private_a = a
                    , private_b = b
                    }
         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
            let m'' = os2ip m'
            return $ i2osp $ expSafe m'' 2 n

-- | 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.
--
-- See algorithm 8.12 in "Handbook of Applied Cryptography" by Alfred J. Menezes et al.
decrypt
    :: HashAlgorithm hash
    => OAEPParams hash ByteString ByteString
    -- ^ OAEP padding parameters
    -> PrivateKey
    -- ^ private key
    -> ByteString
    -- ^ ciphertext
    -> Maybe ByteString
decrypt oaep pk c =
    let p = private_p pk
        q = private_q pk
        a = private_a pk
        b = private_b pk
        n = public_n $ private_pub pk
        k = numBytes n
        c' = os2ip c
        solutions = rights $ toList $ mapTuple (unpad oaep k . i2ospOf_ k) $ sqroot' c' p q a b n
     in case solutions of
            [x] -> Just x
            _ -> Nothing
  where
    toList (w, x, y, z) = w : x : y : z : []
    mapTuple f (w, x, y, z) = (f w, f x, f y, f z)

-- | Sign message using padding, hash algorithm and private key.
--
-- See <https://en.wikipedia.org/wiki/Rabin_signature_algorithm>.
signWith
    :: HashAlgorithm hash
    => ByteString
    -- ^ padding
    -> PrivateKey
    -- ^ private key
    -> hash
    -- ^ hash function
    -> ByteString
    -- ^ message to sign
    -> Either Error Signature
signWith padding pk hashAlg m = do
    h <- calculateHash padding pk hashAlg m
    signature <- calculateSignature h
    return signature
  where
    calculateSignature h =
        let p = private_p pk
            q = private_q pk
            a = private_a pk
            b = private_b pk
            n = public_n $ private_pub pk
         in if h >= n
                then Left MessageTooLong
                else
                    let (r, _, _, _) = sqroot' h p q a b n
                     in Right $ Signature (os2ip padding, r)

-- | Sign message using hash algorithm and private key.
--
-- See <https://en.wikipedia.org/wiki/Rabin_signature_algorithm>.
sign
    :: (MonadRandom m, HashAlgorithm hash)
    => PrivateKey
    -- ^ private key
    -> hash
    -- ^ hash function
    -> ByteString
    -- ^ message to sign
    -> m (Either Error Signature)
sign pk hashAlg m = do
    padding <- findPadding
    return $ signWith padding pk hashAlg m
  where
    findPadding = do
        padding <- getRandomBytes 8
        case calculateHash padding pk hashAlg m of
            Right _ -> return padding
            _ -> findPadding

-- | Calculate hash of message and padding.
-- If the padding is valid, then the result of the hash operation is returned, otherwise an error.
calculateHash
    :: HashAlgorithm hash
    => ByteString
    -- ^ padding
    -> PrivateKey
    -- ^ private key
    -> hash
    -- ^ hash function
    -> ByteString
    -- ^ message to sign
    -> Either Error Integer
calculateHash padding pk hashAlg m =
    let p = private_p pk
        q = private_q pk
        h = os2ip $ hashWith hashAlg $ B.append padding m
     in case (jacobi (h `mod` p) p, jacobi (h `mod` q) q) of
            (Just 1, Just 1) -> Right h
            _ -> Left InvalidParameters

-- | Verify signature using hash algorithm and public key.
--
-- See <https://en.wikipedia.org/wiki/Rabin_signature_algorithm>.
verify
    :: HashAlgorithm hash
    => PublicKey
    -- ^ private key
    -> hash
    -- ^ hash function
    -> ByteString
    -- ^ message
    -> Signature
    -- ^ signature
    -> Bool
verify pk hashAlg m (Signature (padding, s)) =
    let n = public_n pk
        p = i2osp padding
        h = os2ip $ hashWith hashAlg $ B.append p m
        h' = expSafe s 2 n
     in h' == h

-- | Square roots modulo prime p where p is congruent 3 mod 4
-- Value a must be a quadratic residue modulo p (i.e. jacobi symbol (a/n) = 1).
--
-- See algorithm 3.36 in "Handbook of Applied Cryptography" by Alfred J. Menezes et al.
sqroot
    :: Integer
    -> Integer
    -- ^ prime p
    -> (Integer, Integer)
sqroot a p =
    let r = expSafe a ((p + 1) `div` 4) p
     in (r, -r)

-- | Square roots modulo n given its prime factors p and q (both congruent 3 mod 4)
-- Value a must be a quadratic residue of both modulo p and modulo q (i.e. jacobi symbols (a/p) = (a/q) = 1).
--
-- See algorithm 3.44 in "Handbook of Applied Cryptography" by Alfred J. Menezes et al.
sqroot'
    :: Integer
    -> Integer
    -- ^ prime p
    -> Integer
    -- ^ prime q
    -> Integer
    -- ^ c such that c*p + d*q = 1
    -> Integer
    -- ^ d such that c*p + d*q = 1
    -> Integer
    -- ^ n = p*q
    -> (Integer, Integer, Integer, Integer)
sqroot' a p q c d n =
    let (r, _) = sqroot a p
        (s, _) = sqroot a q
        x = (r * d * q + s * c * p) `mod` n
        y = (r * d * q - s * c * p) `mod` n
     in (x, (-x) `mod` n, y, (-y) `mod` n)
