-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

{- | Timelock puzzle algorithms implementation.

__WARNING__: the timelock mechanism described and implemented here is
vulnerable. At the time of writing, no details were released, but creation of
smart contracts using this functionality is disabled since Lima.

This module follows the reference implementation for the most part,
which you can find in the
[tezos repository](https://gitlab.com/tezos/tezos/-/blob/b1a2ff0334405cafd7465bfa991d23844f0b4e70/src/lib_crypto/timelock.ml).

For a more high-level overview of the concepts, refer to the [timelock
documentation page](http://tezos.gitlab.io/011/timelock.html).

The general idea is built upon [Rivest, Shamir, Wagner "Time-lock puzzles and
timed-release Crypto"](http://www.hashcash.org/papers/time-lock.pdf), there
are however some differences from the paper:

* The paper suggests using RC5 cipher, which Tezos implementation
eschews in favor of NaCl's "secret box".
* The paper suggest generating the symmetric secret key \(K\) directly, then
encrypting it with a randomly chosen value \(a\) as \(C_K = K + a^{2^t} \pmod n\).
Tezos implementation instead randomly generates only \(a\), and then produces
the secret key using BLAKE2b KDF with a fixed key from \(a^{2^t} \pmod n\).
* Since the secret key is determined only by the "unlocked" value, the
time-locked value representation also differs. In the paper it's represented as
\((n,a,t,C_K,C_M)\), i.e. the tuple of modulus, "locked" value, time,
encrypted key and encrypted message. In Tezos implementation it's instead
\((a,n,C_M)\), and \(t\) is treated as a separate argument.
* Likely to guard the protocol from guessing attacks, additional "proof"
verification is added, described in
[Boneh, Bünz, Fisch "A Survey of Two Verifiable Delay Functions"](https://eprint.iacr.org/2018/712.pdf)
-}
module Morley.Tezos.Crypto.Timelock
  ( TLTime(.., TLTime)
  , Chest(..)
  , ChestKey(..)
  , Ciphertext(..)
  , OpeningResult(..)
  , createChestAndChestKey
  , createChestKey
  , chestBytes
  , chestKeyBytes
  , chestFromBytes
  , chestKeyFromBytes
  , openChest
  , mkTLTime
  , toTLTime
  -- * Internal, not safe for cryptography
  , createChestAndChestKeyFromSeed
  ) where

import Control.Monad.Random (evalRand, genByteString, getRandomR, liftRand, mkStdGen)
import Crypto.Number.ModArithmetic (expFast)
import Crypto.Number.Prime (findPrimeFrom)
import Crypto.Number.Serialize.LE (i2osp, os2ip)
import Crypto.Sodium.Encrypt.Symmetric qualified as Box (Key, Nonce, decrypt, encrypt)
import Crypto.Sodium.Hash (blake2bWithKey)
import Crypto.Sodium.Nonce qualified as Nonce (generate)
import Crypto.Sodium.Random qualified as Random (generate)
import Data.Binary qualified as Bi (Binary(..), decodeOrFail, encode)
import Data.Binary.Get qualified as Bi (getByteString)
import Data.Binary.Put qualified as Bi (putBuilder)
import Data.Bits (shiftL)
import Data.ByteArray.Sized (SizedByteArray, sizedByteArray, unSizedByteArray)
import Data.ByteString qualified as BS (intercalate, length, replicate)
import Data.ByteString.Lazy qualified as LBS (fromStrict, toStrict)
import Fmt (Buildable(..))
import GHC.TypeNats (Div, type (+))
import Options.Applicative qualified as Opt

import Morley.Micheline.Binary.Internal
import Morley.Util.CLI

-- | RSA-inspired prime factors, which produce (literally) the 'PublicModulus'.
--
-- This value should be kept secret.
data RSAFactors = RSAFactors { rsaP :: Integer, rsaQ :: Integer }
  deriving stock Show

-- | RSA-inspired semi-prime modulus.
newtype PublicModulus = PublicModulus { unPublicModulus :: Integer }
  deriving stock (Show, Eq)
  deriving newtype NFData

-- | The "locked" value. Essentially a random integer between 0 and
-- 'PublicModulus' (exclusive).
newtype Locked = Locked { unLocked :: Integer }
  deriving stock (Show, Eq)
  deriving newtype NFData

-- | The "unlocked" value, i.e. \(a^{2^t} \pmod n\), where \(a\) is 'Locked'
-- value, \(t\) is t'TLTime' and \(n\) is 'PublicModulus'.
newtype Unlocked = Unlocked { unUnlocked :: Integer }
  deriving stock (Show, Eq)
  deriving newtype NFData

-- | The key for the symmetric encryption, \(K\).
newtype SymmetricKey = SymmetricKey (Box.Key ByteString)
  deriving stock Show

{- | A "proof" that the chest was opened fairly, i.e. the key wasn't just
guessed.

The proof is verified by checking that

\[
  a ^ (2 ^ t) = (p ^ l) (a ^ r) \pmod n
\]

which is equivalent to

\[
  2 ^ t = (((2 ^ t) / l) * l) + (2 ^ t \mod l) \pmod {\phi(n)}
\]

where \(a\) is a 'Locked' value, \(t\) is t'TLTime', \(p\) is 'Proof', \(n\) is
'PublicModulus', \(l\) is a prime, chosen deterministically from the hash
of the puzzle, and \(r = 2^t \pmod l\)

What this essentially boils down to, is that we can compute the "proof" either
as

\[
  p = a^{2^t / l \mod {\phi(n)}} \pmod n
\]

if we know the modulo factorization, or

\[
  p = a^{2^t / l} \pmod n
\]

if we don't.

See https://eprint.iacr.org/2018/712.pdf section 3.2.
-}
newtype Proof = Proof { unProof :: Integer }
  deriving stock (Show, Eq)
  deriving newtype NFData

-- | Number of steps a timelock needs to be opened without knowing a "secret",
-- i.e. modulo factorization.
--
-- The reference implementation uses OCaml @int@, and it can only be positive,
-- so on 64-bit architecture it's actually a 62-bit natural. We use 'Word62'
-- to represent it.
--
-- The constructor is marked "Unsafe" since GHC does not warn on overflowing
-- literals (exceeding custom 'Word62' type bounds), thus the resultant
-- t'TLTime' value may get truncated silently.
--
-- >>> UnsafeTLTime 4611686018427387906
-- UnsafeTLTime {unTLTime = 2}
newtype TLTime = UnsafeTLTime { unTLTime :: Word62 }
  deriving stock (Show, Eq)
  deriving newtype Bounded

pattern TLTime :: Word62 -> TLTime
pattern TLTime x <- UnsafeTLTime x
{-# COMPLETE TLTime #-}

instance HasCLReader TLTime where
  getReader = either (readerError . toString) pure . mkTLTime @Word64 =<< Opt.auto
  getMetavar = "TIME"

instance Buildable TLTime where
  build = show . unTLTime

-- | Safely creates t'TLTime' checking for
-- overflow and underflow. Accepts a number of any type.
mkTLTime :: Integral i => i -> Either Text TLTime
mkTLTime = bimap (fromString . displayException) UnsafeTLTime . fromIntegralNoOverflow

-- | Safely creates t'TLTime'.
--
-- This is the recommended way to create t'TLTime' values.
--
-- When constructing literals, you'll need to specify the type of the literal.
-- Bear in mind that GHC will check for literal overflow on builtin types like
-- 'Word16' and 'Word32', but not on 'Word62', so be aware that 'toTLTime' from
-- 'Word62' will overflow silently. Prefer using builtin types when possible.
--
-- >>> unTLTime $ toTLTime (4611686018427387903 :: Word62)
-- 4611686018427387903
-- >>> unTLTime $ toTLTime (4611686018427387904 :: Word62)
-- 0
toTLTime :: (Integral a, CheckIntSubType a Word62) => a -> TLTime
toTLTime = UnsafeTLTime . fromIntegral

-- | A nonce for symmetric encryption.
newtype Nonce = Nonce { unNonce :: Box.Nonce ByteString }
  deriving stock (Show, Eq)

instance NFData Nonce where
  rnf (Nonce x) = rnf $ unSizedByteArray x

-- | Ciphertext with nonce.
data Ciphertext = Ciphertext
  { ctNonce :: Nonce
  , ctPayload :: ByteString
  } deriving stock (Show, Eq, Generic)
    deriving anyclass NFData

instance Bi.Binary Ciphertext where
  put Ciphertext{..} = Bi.putBuilder $
      buildByteString (unSizedByteArray $ unNonce ctNonce)
   <> buildDynamic buildByteString (DynamicSize ctPayload)
  get = do
    mbNonce <- sizedByteArray <$> Bi.getByteString 24
    ctNonce <- case mbNonce of
      Just sza -> pure $ Nonce sza
      Nothing -> fail "Incorrect nonce size"
        -- NB: this shouldn't happen unless NaCl box changes nonce size
    DynamicSize ctPayload <- getDynamic getByteString
    when (BS.length ctPayload <= secretBoxTagBytes) $
      fail "Ciphertext size <= 0"
    pure $ Ciphertext{..}
    where
      -- This is hard-coded in the reference implementation,
      -- https://gitlab.com/tezos/tezos/-/blob/b1a2ff0334405cafd7465bfa991d23844f0b4e70/src/lib_hacl_glue/unix/hacl.ml#L183
      -- but essentially this is the length of a NaCl box ciphertext with
      -- empty payload
      secretBoxTagBytes = 16

-- | A chest "key" with proof that it was indeed opened fairly.
data ChestKey = ChestKey
  { ckUnlockedVal :: Unlocked
  , ckProof :: Proof
  } deriving stock (Show, Eq, Generic)
    deriving anyclass NFData

instance Bi.Binary ChestKey where
  put ChestKey{..} = Bi.putBuilder $
    buildNatural (unUnlocked ckUnlockedVal) <> buildNatural (unProof ckProof)
  get = ChestKey <$> (Unlocked <$> getNatural) <*> (Proof <$> getNatural)

-- | A locked chest
data Chest = Chest
  { chestLockedVal :: Locked
  , chestPublicModulus :: PublicModulus
  , chestCiphertext :: Ciphertext
  } deriving stock (Show, Eq, Generic)
    deriving anyclass NFData

instance Bi.Binary Chest where
  put Chest{..} = do
    Bi.putBuilder $ buildNatural (unLocked chestLockedVal)
                 <> buildNatural (unPublicModulus chestPublicModulus)
    Bi.put chestCiphertext

  get = do
    chestLockedVal <- Locked <$> getNatural
    chestPublicModulus <- PublicModulus <$> getNatural
    when (unLocked chestLockedVal >= unPublicModulus chestPublicModulus) $
      fail "locked value is not in the rsa group"
    when (unPublicModulus chestPublicModulus <= minPublicModulus) $
      fail "public modulus is too small"
    chestCiphertext <- Bi.get
    pure Chest{..}
    where
      -- hard-coded in the reference implementation
      -- https://gitlab.com/tezos/tezos/-/blob/b1a2ff0334405cafd7465bfa991d23844f0b4e70/src/lib_crypto/timelock.ml#L193
      minPublicModulus = shiftL 2 2000

type SizeModulus = 256 -- bytes, i.e. 2048 bits
type HalfModulus = Div SizeModulus 2 -- bytes, i.e. 1024 bits

randomInt :: forall n. (KnownNat n) => IO Integer
randomInt = os2ip <$> Random.generate @ByteString @n

randomPrime :: forall n. (KnownNat n) => IO Integer
randomPrime = findPrimeFrom <$> randomInt @n

unlockedValueToSymmetricKey :: Unlocked -> SymmetricKey
unlockedValueToSymmetricKey (Unlocked value) =
  -- "Tezoskdftimelockv0" is hard-coded in the reference implementation
  -- see https://gitlab.com/tezos/tezos/-/blob/b1a2ff0334405cafd7465bfa991d23844f0b4e70/src/lib_crypto/timelock.ml#L47
  let key :: ByteString = "Tezoskdftimelockv0"
      str :: ByteString = show value
  in SymmetricKey $ blake2bWithKey @32 @ByteString key str

genRSAFactors :: IO (PublicModulus, RSAFactors)
genRSAFactors = do
  p <- randomPrime @HalfModulus
  q <- randomPrime @HalfModulus
  pure (PublicModulus (p * q), RSAFactors p q)

genLockedValue :: PublicModulus -> IO Locked
genLockedValue (PublicModulus pub) = do
  z <- randomInt @(SizeModulus + 16)
  pure . Locked $ z `mod` pub

hashToPrime :: PublicModulus -> TLTime -> Locked -> Unlocked -> Integer
hashToPrime (PublicModulus pub) (TLTime time) (Locked locked) (Unlocked unlocked) =
  -- "\32" and "\xff\x00\xff\x00\xff\x00\xff\x00" are hard-coded in the reference implementation
  -- see https://gitlab.com/tezos/tezos/-/blob/b1a2ff0334405cafd7465bfa991d23844f0b4e70/src/lib_crypto/timelock.ml#L78
  -- and https://gitlab.com/tezos/tezos/-/blob/b1a2ff0334405cafd7465bfa991d23844f0b4e70/src/lib_crypto/timelock.ml#L81
  let personalization :: ByteString = "\32"
      s = BS.intercalate "\xff\x00\xff\x00\xff\x00\xff\x00" $
        show time : map (pad . i2osp) [pub, locked, unlocked]
      hash_result :: SizedByteArray 32 ByteString = blake2bWithKey personalization s
  in findPrimeFrom (os2ip hash_result)
  where
    -- pads right with zero bytes so that length is multiple of 8
    -- this is needed due to a quirk of how @Z.to_bits@ works in OCaml
    pad bs =
      let len = length bs
          newlen = ceiling ((fromIntegralOverflowing len :: Rational) / 8) * 8
          diff = newlen - len
      in bs <> BS.replicate diff 0

proveWithoutSecret :: PublicModulus -> TLTime -> Locked -> Unlocked -> Proof
proveWithoutSecret (PublicModulus pub) time (Locked locked) (Unlocked unlocked) =
  let l = hashToPrime (PublicModulus pub) time (Locked locked) (Unlocked unlocked)
      pow = (2 ^ unTLTime time) `div` l
  in Proof $ expFast locked pow pub

verifyTimeLock :: PublicModulus -> TLTime -> Locked -> Unlocked -> Proof -> Bool
verifyTimeLock (PublicModulus pub) time (Locked locked) (Unlocked unlocked) (Proof proof) =
  let l = hashToPrime (PublicModulus pub) time (Locked locked) (Unlocked unlocked)
      r = expFast 2 (toInteger $ unTLTime time) l
  in unlocked == (expFast proof l pub * expFast locked r pub) `mod` pub

unlockAndProveWithSecret :: RSAFactors -> TLTime -> Locked -> (Unlocked, Proof)
unlockAndProveWithSecret RSAFactors{..} time (Locked locked) =
  (Unlocked unlocked, Proof proof)
  where
    phi = (rsaP - 1) * (rsaQ - 1)
    pub = rsaP * rsaQ
    e = expFast 2 (toInteger $ unTLTime time) phi
    unlocked = expFast locked e pub
    l = hashToPrime (PublicModulus pub) time (Locked locked) (Unlocked unlocked)
    pow = ((2 ^ unTLTime time) `div` l) `mod` phi
    proof = expFast locked pow pub

unlockWithoutSecret :: PublicModulus -> TLTime -> Locked -> Unlocked
unlockWithoutSecret (PublicModulus pub) (TLTime time) (Locked locked) =
  Unlocked $ go time locked
  where
    go 0 v = v
    go t v = go (pred t) (v * v `mod` pub)

encrypt :: SymmetricKey -> ByteString -> IO Ciphertext
encrypt (SymmetricKey key) plaintext = do
  nonce <- Nonce.generate
  let payload = Box.encrypt key nonce plaintext
  pure $ Ciphertext (Nonce nonce) payload

decrypt :: SymmetricKey -> Ciphertext -> Maybe ByteString
decrypt (SymmetricKey key) Ciphertext{ctNonce = Nonce nonce, ..}
  = Box.decrypt key nonce ctPayload

-- | Create a timelock puzzle and a key.
createChestAndChestKey
  :: ByteString -- ^ Chest content
  -> TLTime -- ^ Time (in elementary actions) to open without key.
  -> IO (Chest, ChestKey)
createChestAndChestKey payload time = do
  (pub, secret) <- genRSAFactors
  locked <- genLockedValue pub
  let (unlocked, proof) = unlockAndProveWithSecret secret time locked
      key = unlockedValueToSymmetricKey unlocked
  ciphertext <- encrypt key payload
  pure $ (Chest locked pub ciphertext, ChestKey unlocked proof)

-- | Forge a chest key the hard way.
createChestKey :: Chest -> TLTime -> ChestKey
createChestKey Chest{..} time =
  let unlocked = unlockWithoutSecret chestPublicModulus time chestLockedVal
      proof = proveWithoutSecret chestPublicModulus time chestLockedVal unlocked
  in ChestKey unlocked proof

-- | The result of opening the chest.
data OpeningResult
  = Correct ByteString -- ^ The chest was opened correctly.
  | BogusCipher -- ^ The chest was opened correctly, but the contents do not decode with
  -- the given symmetric key.
  | BogusOpening -- ^ The chest was not opened correctly, i.e. proof verification failed.
  deriving stock (Show, Eq)

-- | Try to (quickly) open a chest with the given key, verifying the proof.
openChest :: Chest -> ChestKey -> TLTime -> OpeningResult
openChest Chest{..} ChestKey{..} time
  | verifyTimeLock chestPublicModulus time chestLockedVal ckUnlockedVal ckProof
  = maybe BogusCipher Correct $
      decrypt (unlockedValueToSymmetricKey ckUnlockedVal) chestCiphertext
  | otherwise = BogusOpening

-- | Convert a 'ChestKey' to binary representation, used by Tezos
chestKeyBytes :: ChestKey -> ByteString
chestKeyBytes = LBS.toStrict . Bi.encode

-- | Convert a 'Chest' to binary representation, used by Tezos
chestBytes :: Chest -> ByteString
chestBytes = LBS.toStrict . Bi.encode

-- | Read a 'Chest' from binary representation, used by Tezos
chestFromBytes :: ByteString -> Either Text Chest
chestFromBytes bs = case Bi.decodeOrFail $ LBS.fromStrict bs of
  Right (trail, _, res)
    | null trail -> Right res
    | otherwise -> Left "trailing unconsumed bytes"
  Left (_, _, err) -> Left (fromString err)

-- | Read a 'ChestKey' from binary representation, used by Tezos
chestKeyFromBytes :: ByteString -> Either Text ChestKey
chestKeyFromBytes bs = case Bi.decodeOrFail $ LBS.fromStrict bs of
  Right (trail, _, res)
    | null trail -> Right res
    | otherwise -> Left "trailing unconsumed bytes"
  Left (_, _, err) -> Left (fromString err)

-- | Construct a chest purely based on a seed for pseudorandom generator.
-- This is not suitable for cryptography, used in tests.
createChestAndChestKeyFromSeed
  :: Int -- ^ Pseudo-random seed
  -> ByteString -- ^ Chest content
  -> TLTime -- ^ TLTime (in elementary actions) to open without key.
  -> (Chest, ChestKey)
createChestAndChestKeyFromSeed seed payload time = flip evalRand (mkStdGen seed) do
  let rangeLow = 2 ^ (natVal @HalfModulus Proxy * 8)
      rangeHigh = rangeLow * 2 - 1
      range = (rangeLow, rangeHigh)
  p' <- getRandomR range
  q' <- getRandomR range
  let pub = PublicModulus pub'
      pub' = p * q
      p = findPrimeFrom p'
      q = findPrimeFrom q'
      secret = RSAFactors p q
      lockedMax = 2 ^ (natVal @SizeModulus Proxy + 16) * 8 - 1
  locked <- Locked . (`mod` pub') <$> getRandomR (0, lockedMax)
  let (unlocked, proof) = unlockAndProveWithSecret secret time locked
      SymmetricKey key = unlockedValueToSymmetricKey unlocked
  nonce <- fromMaybe (error "impossible") . sizedByteArray <$> liftRand (genByteString 24)
  let ciphertext = Ciphertext (Nonce nonce) $ Box.encrypt key nonce payload
  pure $ (Chest locked pub ciphertext, ChestKey unlocked proof)