{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE UndecidableInstances #-} -- for Reifies instances
{-# OPTIONS_GHC -fno-warn-orphans #-}
-- | Finite Field Cryptography (FFC)
-- is a method of implementing discrete logarithm cryptography
-- using finite field mathematics.
module Voting.Protocol.FFC
 ( module Voting.Protocol.FFC
 , Natural
 , Random.RandomGen
 , Reifies(..), reify
 , Proxy(..)
 ) where

import Control.Arrow (first)
import Control.DeepSeq (NFData)
import Control.Monad (Monad(..), unless)
import Data.Aeson (ToJSON(..),FromJSON(..),(.:),(.:?),(.=))
import Data.Bits
import Data.Bool
import Data.Either (Either(..))
import Data.Eq (Eq(..))
import Data.Foldable (Foldable, foldl')
import Data.Function (($), (.), id)
import Data.Functor ((<$>))
import Data.Int (Int)
import Data.Maybe (Maybe(..), fromMaybe, fromJust)
import Data.Monoid (Monoid(..))
import Data.Ord (Ord(..))
import Data.Proxy (Proxy(..))
import Data.Reflection (Reifies(..), reify)
import Data.Semigroup (Semigroup(..))
import Data.String (IsString(..))
import Data.Text (Text)
import GHC.Generics (Generic)
import GHC.Natural (minusNaturalMaybe)
import Numeric.Natural (Natural)
import Prelude (Integer, Integral(..), fromIntegral, Enum(..))
import Text.Read (readMaybe, readEither)
import Text.Show (Show(..))
import qualified Control.Monad.Trans.State.Strict as S
import qualified Crypto.Hash as Crypto
import qualified Data.Aeson as JSON
import qualified Data.Aeson.Types as JSON
import qualified Data.ByteArray as ByteArray
import qualified Data.ByteString as BS
import qualified Data.ByteString.Base64 as BS64
import qualified Data.Char as Char
import qualified Data.List as List
import qualified Data.Text as Text
import qualified Data.Text.Encoding as Text
import qualified Data.Text.Lazy as TL
import qualified Data.Text.Lazy.Builder as TLB
import qualified Data.Text.Lazy.Builder.Int as TLB
import qualified Prelude as Num
import qualified System.Random as Random

-- * Type 'FFC'
-- | Mutiplicative Sub-Group of a Finite Prime Field.
--
-- NOTE: an 'FFC' term-value is brought into the context of many functions
-- through a type-variable @c@ whose 'Reifies' constraint enables to 'reflect'
-- that 'FFC' at the term-level (a surprising technique but a very useful one).
-- Doing like this is simpler than working in a 'Monad' (like a 'Reader'),
-- and enables that 'FFC' term to be used simply in instances' methods
-- not supporting an inner 'Monad', like 'parseJSON', 'randomR', 'fromEnum' or 'arbitrary'.
-- Aside from that, the sharing of 'FFC' amongst several types
-- is encoded at the type-level by including @c@
-- as a phantom type of 'F', 'G' and 'E'.
data FFC = FFC
 {   ffc_name :: Text
 ,   ffc_fieldCharac :: !Natural
     -- ^ The prime number characteristic of a Finite Prime Field.
     --
     -- ElGamal's hardness to decrypt requires a large prime number
     -- to form the 'Multiplicative' subgroup.
 ,   ffc_groupGen :: !Natural
     -- ^ A generator of the 'Multiplicative' subgroup of the Finite Prime Field.
     --
     -- NOTE: since 'ffc_fieldCharac' is prime,
     -- the 'Multiplicative' subgroup is cyclic,
     -- and there are phi('fieldCharac'-1) many choices for the generator of the group,
     -- where phi is the Euler totient function.
 ,   ffc_groupOrder :: !Natural
     -- ^ The order of the subgroup.
     --
     -- WARNING: 'ffc_groupOrder' MUST be a prime number dividing @('ffc_fieldCharac'-1)@
     -- to ensure that ElGamal is secure in terms of the DDH assumption.
 } deriving (Eq,Show,Generic,NFData)
instance ToJSON FFC where
        toJSON FFC{..} =
                JSON.object $
                 (if Text.null ffc_name then [] else ["name" .= ffc_name] ) <>
                 [ "p" .= show ffc_fieldCharac
                 , "g" .= show ffc_groupGen
                 , "q" .= show ffc_groupOrder
                 ]
        toEncoding FFC{..} =
                JSON.pairs $
                        (if Text.null ffc_name then mempty else "name" .= ffc_name) <>
                        "p" .= show ffc_fieldCharac <>
                        "g" .= show ffc_groupGen <>
                        "q" .= show ffc_groupOrder
instance FromJSON FFC where
        parseJSON = JSON.withObject "FFC" $ \o -> do
                ffc_name <- fromMaybe "" <$> (o .:? "name")
                p <- o .: "p"
                g <- o .: "g"
                q <- o .: "q"
                -- TODO: check p is probable prime
                -- TODO: check q is probable prime
                ffc_fieldCharac <- case readEither (Text.unpack p) of
                 Left err -> JSON.typeMismatch ("FFC: fieldCharac: "<>err) (JSON.String p)
                 Right a -> return a
                ffc_groupGen <- case readEither (Text.unpack g) of
                 Left err -> JSON.typeMismatch ("FFC: groupGen: "<>err) (JSON.String g)
                 Right a -> return a
                ffc_groupOrder <- case readEither (Text.unpack q) of
                 Left err -> JSON.typeMismatch ("FFC: groupOrder: "<>err) (JSON.String q)
                 Right a -> return a
                unless (nat ffc_groupGen < ffc_fieldCharac) $
                        JSON.typeMismatch "FFC: groupGen is not lower than fieldCharac" (JSON.Object o)
                unless (ffc_groupOrder < ffc_fieldCharac) $
                        JSON.typeMismatch "FFC: groupOrder is not lower than fieldCharac" (JSON.Object o)
                unless (nat ffc_groupGen > 1) $
                        JSON.typeMismatch "FFC: groupGen is not greater than 1" (JSON.Object o)
                unless (fromJust (ffc_fieldCharac`minusNaturalMaybe`one) `rem` ffc_groupOrder == 0) $
                        JSON.typeMismatch "FFC: groupOrder does not divide fieldCharac-1" (JSON.Object o)
                return FFC{..}

fieldCharac :: forall c. Reifies c FFC => Natural
fieldCharac = ffc_fieldCharac (reflect (Proxy::Proxy c))

groupGen :: forall c. Reifies c FFC => G c
groupGen = G $ F $ ffc_groupGen (reflect (Proxy::Proxy c))

groupOrder :: forall c. Reifies c FFC => Natural
groupOrder = ffc_groupOrder (reflect (Proxy::Proxy c))

-- ** Examples
-- | Weak parameters for debugging purposes only.
weakFFC :: FFC
weakFFC = FFC
 { ffc_name        = "weakFFC"
 , ffc_fieldCharac = 263
 , ffc_groupGen    = 2
 , ffc_groupOrder  = 131
 }

-- | Parameters used in Belenios.
-- A 2048-bit 'fieldCharac' of a Finite Prime Field,
-- with a 256-bit 'groupOrder' for a 'Multiplicative' subgroup
-- generated by 'groupGen'.
beleniosFFC :: FFC
beleniosFFC = FFC
 { ffc_name        = "beleniosFFC"
 , ffc_fieldCharac = 20694785691422546401013643657505008064922989295751104097100884787057374219242717401922237254497684338129066633138078958404960054389636289796393038773905722803605973749427671376777618898589872735865049081167099310535867780980030790491654063777173764198678527273474476341835600035698305193144284561701911000786737307333564123971732897913240474578834468260652327974647951137672658693582180046317922073668860052627186363386088796882120769432366149491002923444346373222145884100586421050242120365433561201320481118852408731077014151666200162313177169372189248078507711827842317498073276598828825169183103125680162072880719
 , ffc_groupGen    =  2402352677501852209227687703532399932712287657378364916510075318787663274146353219320285676155269678799694668298749389095083896573425601900601068477164491735474137283104610458681314511781646755400527402889846139864532661215055797097162016168270312886432456663834863635782106154918419982534315189740658186868651151358576410138882215396016043228843603930989333662772848406593138406010231675095763777982665103606822406635076697764025346253773085133173495194248967754052573659049492477631475991575198775177711481490920456600205478127054728238140972518639858334115700568353695553423781475582491896050296680037745308460627
 , ffc_groupOrder  = 78571733251071885079927659812671450121821421258408794611510081919805623223441
 }

-- * Type 'F'
-- | The type of the elements of a Finite Prime Field.
--
-- A field must satisfy the following properties:
--
-- * @(f, ('+'), 'zero')@ forms an abelian group,
--   called the 'Additive' group of 'f'.
--
-- * @('NonNull' f, ('*'), 'one')@ forms an abelian group,
--   called the 'Multiplicative' group of 'f'.
--
-- * ('*') is associative:
--   @(a'*'b)'*'c == a'*'(b'*'c)@ and
--   @a'*'(b'*'c) == (a'*'b)'*'c@.
--
-- * ('*') and ('+') are both commutative:
--   @a'*'b == b'*'a@ and
--   @a'+'b == b'+'a@
--
-- * ('*') and ('+') are both left and right distributive:
--   @a'*'(b'+'c) == (a'*'b) '+' (a'*'c)@ and
--   @(a'+'b)'*'c == (a'*'c) '+' (b'*'c)@
--
-- The 'Natural' is always within @[0..'fieldCharac'-1]@.
newtype F c = F { unF :: Natural }
 deriving (Eq,Ord,Show)
 deriving newtype NFData
instance ToJSON (F c) where
        toJSON (F x) = JSON.toJSON (show x)
instance Reifies c FFC => FromJSON (F c) where
        parseJSON (JSON.String s)
         | Just (c0,_) <- Text.uncons s
         , c0 /= '0'
         , Text.all Char.isDigit s
         , Just x <- readMaybe (Text.unpack s)
         , x < fieldCharac @c
         = return (F x)
        parseJSON json = JSON.typeMismatch "FieldElement" json
instance Reifies c FFC => FromNatural (F c) where
        fromNatural i = F $ abs $ i `mod` fieldCharac @c
                where
                abs x | x < 0 = x + fieldCharac @c
                      | otherwise = x
instance ToNatural (F c) where
        nat = unF
instance Reifies c FFC => Additive (F c) where
        zero = F 0
        F x + F y = F $ (x + y) `mod` fieldCharac @c
instance Reifies c FFC => Negable (F c) where
        neg (F x)
         | x == 0 = zero
         | otherwise = F $ fromJust $ nat (fieldCharac @c)`minusNaturalMaybe`x
instance Reifies c FFC => Multiplicative (F c) where
        one = F 1
        F x * F y = F $ (x * y) `mod` fieldCharac @c
instance Reifies c FFC => Random.Random (F c) where
        randomR (F lo, F hi) =
                first (F . fromIntegral) .
                Random.randomR
                 ( 0`max`toInteger lo
                 , toInteger hi`min`(toInteger (fieldCharac @c) - 1) )
        random =
                first (F . fromIntegral) .
                Random.randomR (0, toInteger (fieldCharac @c) - 1)

-- ** Class 'Additive'
class Additive a where
        zero :: a
        (+) :: a -> a -> a; infixl 6 +
        sum :: Foldable f => f a -> a
        sum = foldl' (+) zero
instance Additive Natural where
        zero = 0
        (+)  = (Num.+)
instance Additive Integer where
        zero = 0
        (+)  = (Num.+)
instance Additive Int where
        zero = 0
        (+)  = (Num.+)

-- *** Class 'Negable'
class Additive a => Negable a where
        neg :: a -> a
        (-) :: a -> a -> a; infixl 6 -
        x-y = x + neg y
instance Negable Integer where
        neg  = Num.negate
instance Negable Int where
        neg  = Num.negate

-- ** Class 'Multiplicative'
class Multiplicative a where
        one :: a
        (*) :: a -> a -> a; infixl 7 *
instance Multiplicative Natural where
        one = 1
        (*) = (Num.*)
instance Multiplicative Integer where
        one = 1
        (*) = (Num.*)
instance Multiplicative Int where
        one = 1
        (*) = (Num.*)

-- ** Class 'Invertible'
class Multiplicative a => Invertible a where
        inv :: a -> a
        (/) :: a -> a -> a; infixl 7 /
        x/y = x * inv y

-- * Type 'G'
-- | The type of the elements of a 'Multiplicative' subgroup of a Finite Prime Field.
newtype G c = G { unG :: F c }
 deriving (Eq,Ord,Show)
 deriving newtype NFData
instance ToJSON (G c) where
        toJSON (G x) = JSON.toJSON x
instance Reifies c FFC => FromJSON (G c) where
        parseJSON (JSON.String s)
         | Just (c0,_) <- Text.uncons s
         , c0 /= '0'
         , Text.all Char.isDigit s
         , Just x <- readMaybe (Text.unpack s)
         , x < fieldCharac @c
         , r <- G (F x)
         , r ^ E (groupOrder @c) == one
         = return r
        parseJSON json = JSON.typeMismatch "GroupElement" json
instance Reifies c FFC => FromNatural (G c) where
        fromNatural = G . fromNatural
instance ToNatural (G c) where
        nat = unF . unG
instance Reifies c FFC => Multiplicative (G c) where
        one = G $ F one
        G x * G y = G (x * y)
instance Reifies c FFC => Invertible (G c) where
        -- | NOTE: add 'groupOrder' so the exponent given to (^) is positive.
        inv = (^ E (fromJust $ groupOrder @c`minusNaturalMaybe`1))

-- | 'groupGenInverses' returns the infinite list
-- of 'inv'erse powers of 'groupGen':
-- @['groupGen' '^' 'neg' i | i <- [0..]]@,
-- but by computing each value from the previous one.
--
-- Used by 'intervalDisjunctions'.
groupGenInverses :: forall c. Reifies c FFC => [G c]
groupGenInverses = go one
        where
        invGen = inv $ groupGen @c
        go g = g : go (g * invGen)

groupGenPowers :: forall c. Reifies c FFC => [G c]
groupGenPowers = go one
        where go g = g : go (g * groupGen @c)

-- ** Type 'Hash'
newtype Hash c = Hash (E c)
 deriving (Eq,Ord,Show)
 deriving newtype NFData

-- | @('hash' bs gs)@ returns as a number in 'E'
-- the 'Crypto.SHA256' hash of the given 'BS.ByteString' 'bs'
-- prefixing the decimal representation of given subgroup elements 'gs',
-- with a comma (",") intercalated between them.
--
-- NOTE: to avoid any collision when the 'hash' function is used in different contexts,
-- a message 'gs' is actually prefixed by a 'bs' indicating the context.
--
-- Used by 'proveEncryption' and 'verifyEncryption',
-- where the 'bs' usually contains the 'statement' to be proven,
-- and the 'gs' contains the 'commitments'.
hash :: Reifies c FFC => BS.ByteString -> [G c] -> E c
hash bs gs = do
        let s = bs <> BS.intercalate (fromString ",") (bytesNat <$> gs)
        let h = Crypto.hashWith Crypto.SHA256 s
        fromNatural $
                decodeBigEndian $ ByteArray.convert h

-- | @('decodeBigEndian' bs)@ interpret @bs@ as big-endian number.
decodeBigEndian :: BS.ByteString -> Natural
decodeBigEndian =
        BS.foldl'
         (\acc b -> acc`shiftL`8 + fromIntegral b)
         (0::Natural)

-- ** Type 'Base64SHA256'
newtype Base64SHA256 = Base64SHA256 Text
 deriving (Eq,Ord,Show,Generic)
 deriving anyclass (ToJSON,FromJSON)
 deriving newtype NFData

-- | @('base64SHA256' bs)@ returns the 'Crypto.SHA256' hash
-- of the given 'BS.ByteString' 'bs',
-- as a 'Text' escaped in @base64@ encoding
-- (<https://tools.ietf.org/html/rfc4648 RFC 4648>).
base64SHA256 :: BS.ByteString -> Base64SHA256
base64SHA256 bs =
        let h = Crypto.hashWith Crypto.SHA256 bs in
        Base64SHA256 $
                Text.takeWhile (/= '=') $ -- NOTE: no padding.
                Text.decodeUtf8 $ BS64.encode $ ByteArray.convert h

-- ** Type 'HexSHA256'
newtype HexSHA256 = HexSHA256 Text
 deriving (Eq,Ord,Show,Generic)
 deriving anyclass (ToJSON,FromJSON)
 deriving newtype NFData
-- | @('hexSHA256' bs)@ returns the 'Crypto.SHA256' hash
-- of the given 'BS.ByteString' 'bs', escaped in hexadecimal
-- into a 'Text' of 32 lowercase characters.
--
-- Used (in retro-dependencies of this library) to hash
-- the 'PublicKey' of a voter or a trustee.
hexSHA256 :: BS.ByteString -> Text
hexSHA256 bs =
        let h = Crypto.hashWith Crypto.SHA256 bs in
        let n = decodeBigEndian $ ByteArray.convert h in
        -- NOTE: always set the 256 bit then remove it
        -- to always have leading zeros,
        -- and thus always 64 characters wide hashes.
        TL.toStrict $
        TL.tail $ TLB.toLazyText $ TLB.hexadecimal $
        setBit n 256

-- * Type 'E'
-- | An exponent of a (necessarily cyclic) subgroup of a Finite Prime Field.
-- The value is always in @[0..'groupOrder'-1]@.
newtype E c = E { unE :: Natural }
 deriving (Eq,Ord,Show)
 deriving newtype NFData
instance ToJSON (E c) where
        toJSON (E x) = JSON.toJSON (show x)
instance Reifies c FFC => FromJSON (E c) where
        parseJSON (JSON.String s)
         | Just (c0,_) <- Text.uncons s
         , c0 /= '0'
         , Text.all Char.isDigit s
         , Just x <- readMaybe (Text.unpack s)
         , x < groupOrder @c
         = return (E x)
        parseJSON json = JSON.typeMismatch "Exponent" json

instance Reifies c FFC => FromNatural (E c) where
        fromNatural i =
                E $ abs $ i `mod` groupOrder @c
                where
                abs x | x < 0 = x + groupOrder @c
                      | otherwise = x
instance ToNatural (E c) where
        nat = unE

instance Reifies c FFC => Additive (E c) where
        zero = E zero
        E x + E y = E $ (x + y) `mod` groupOrder @c
instance Reifies c FFC => Negable (E c) where
        neg (E x)
         | x == 0 = zero
         | otherwise = E $ fromJust $ nat (groupOrder @c)`minusNaturalMaybe`x
instance Reifies c FFC => Multiplicative (E c) where
        one = E one
        E x * E y = E $ (x * y) `mod` groupOrder @c
instance Reifies c FFC => Random.Random (E c) where
        randomR (E lo, E hi) =
                first (E . fromIntegral) .
                Random.randomR
                 ( 0`max`toInteger lo
                 , toInteger hi`min`(toInteger (groupOrder @c) - 1) )
        random =
                first (E . fromIntegral) .
                Random.randomR (0, toInteger (groupOrder @c) - 1)
instance Reifies c FFC => Enum (E c) where
        toEnum = fromNatural . fromIntegral
        fromEnum = fromIntegral . nat
        enumFromTo lo hi = List.unfoldr
         (\i -> if i<=hi then Just (i, i+one) else Nothing) lo

infixr 8 ^
-- | @(b '^' e)@ returns the modular exponentiation of base 'b' by exponent 'e'.
(^) :: Reifies c FFC => G c -> E c -> G c
(^) b (E e)
 | e == 0 = one
 | otherwise = t * (b*b) ^ E (e`shiftR`1)
        where
        t | testBit e 0 = b
          | otherwise   = one

-- | @('randomR' i)@ returns a random integer in @[0..i-1]@.
randomR ::
 Monad m =>
 Random.RandomGen r =>
 Random.Random i =>
 Negable i =>
 Multiplicative i =>
 i -> S.StateT r m i
randomR i = S.StateT $ return . Random.randomR (zero, i-one)

-- | @('random')@ returns a random integer
-- in the range determined by its type.
random ::
 Monad m =>
 Random.RandomGen r =>
 Random.Random i =>
 Negable i =>
 Multiplicative i =>
 S.StateT r m i
random = S.StateT $ return . Random.random

instance Random.Random Natural where
        randomR (mini,maxi) =
                first (fromIntegral::Integer -> Natural) .
                Random.randomR (fromIntegral mini, fromIntegral maxi)
        random = first (fromIntegral::Integer -> Natural) . Random.random

-- * Conversions

-- ** Class 'FromNatural'
class FromNatural a where
        fromNatural :: Natural -> a

-- ** Class 'ToNatural'
class ToNatural a where
        nat :: a -> Natural
instance ToNatural Natural where
        nat = id

-- | @('bytesNat' x)@ returns the serialization of 'x'.
bytesNat :: ToNatural n => n -> BS.ByteString
bytesNat = fromString . show . nat