{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}

-- | Super-high-level cryptography needed for our DID-inspired implementation.
module DoAuth.Crypto
  ( PKC (..),
    -- | For testing purposes
    deriveKeypair,
    dummyPkc,
    -- | Need to export DummySigned for testing purposes, since it's impossible
    -- to create an invalid signature with just dummyPkc.
    DummySigned,
    sodiumPkc,
    Signature (..),
    kdf,
    rekdf,
    sign,
    verify,
  )
where

import Control.Monad (void)
import Control.Monad.IO.Class (MonadIO (..))
import Crypto.Key (DerivationSlip, Params (..), derive, rederive)
import qualified Crypto.Key as Key
import Crypto.Nonce (generate)
import Crypto.Sign (PublicKey, SecretKey)
import Data.ByteArray (ByteArray, ByteArrayAccess (copyByteArrayToPtr, withByteArray), ScrubbedBytes)
import Data.ByteArray.Sized (ByteArrayN)
import qualified Data.ByteArray.Sized as Sized
import Data.ByteString (ByteString (..))
import Data.Proxy (Proxy (Proxy))
import GHC.TypeLits (type (<=))
import qualified Libsodium as Na
import qualified NaCl.Sign as Sign

-- | We will use only signing primitives, but we are including encryption for
-- completness.
--
-- Ideally, we would like to use RankNTypes here for signature algorithm to work
-- with any blobs, perhaps with some constraint. But for the time being, we'll
-- be working with more rigid construction.
data PKC io pass sk pk sl signature sigdata = PKC
  { -- | Initial key derivation function
    dacpkc_kdf :: pass -> io (Maybe (pk, sk, sl)),
    -- | Rederive with sodiumKdf
    dacpkc_rekdf :: sl -> pass -> Maybe (pk, sk),
    -- | Sign data of type @a@ with key @sk@ and produces a detached @signature@
    -- containing @pk@ for verification.
    dacpkc_sign :: (pk, sk) -> sigdata -> signature pk,
    -- | Verifies @signature@ container's validity.
    dacpkc_verify :: signature pk -> sigdata -> Bool
  }

---------- Sodium crypto -----------

kdf :: ByteString -> IO (Maybe (PublicKey ByteString, SecretKey ScrubbedBytes, DerivationSlip))
kdf = dacpkc_kdf sodiumPkc

rekdf :: DerivationSlip -> ByteString -> Maybe (PublicKey ByteString, SecretKey ScrubbedBytes)
rekdf = dacpkc_rekdf sodiumPkc

sign :: (PublicKey ByteString, SecretKey ScrubbedBytes) -> ByteString -> Signature (PublicKey ByteString)
sign = dacpkc_sign sodiumPkc

verify :: Signature (PublicKey ByteString) -> ByteString -> Bool
verify = dacpkc_verify sodiumPkc

sodiumPkc :: PKC IO ByteString (SecretKey ScrubbedBytes) (PublicKey ByteString) DerivationSlip Signature ByteString
sodiumPkc =
  PKC
    { dacpkc_kdf = sodiumKdf,
      dacpkc_rekdf = sodiumRekdf,
      dacpkc_sign = sodiumSign,
      dacpkc_verify = sodiumVerify
    }

-- | Ops and mem limits are dedicated to making passwords more resistant to
-- brute-force attacks by requiring a considerable amount of resources to
-- produce a hash of said password.
--
-- At the moment we're going with moderate params (256 MiB of RAM and under 1
-- CPU second at 2.8 GHz), but to be more democratic, we'll need to:
--  1. Make the params flexible to acommodate for smaller, slower hardware
--  2. Make sure that clients are ruthlessly denying bad passwords
--
-- I decided to write down integers manually, instead of converting those
-- from CSize, because I'm lazy.
moderateParams :: Params
moderateParams =
  Params
    { opsLimit = 3, -- Na.crypto_pwhash_opslimit_moderate, -- == 3
      memLimit = 268435456 -- Na.crypto_pwhash_memlimit_moderate  -- == 268435456
    }

-- | PKC-compatible deriveKeypair.
--
-- Note: this function is not thread-safe (since the underlying C function is
-- not thread-safe both in Sodium and in NaCl)! Either make sure
-- there are no concurrent calls or see @Crypto.Init@ in
-- <https://hackage.haskell.org/package/crypto-sodium crypto-sodium> to learn
-- how to make this function thread-safe.
sodiumKdf :: ByteString -> IO (Maybe (PublicKey ByteString, SecretKey ScrubbedBytes, DerivationSlip))
sodiumKdf = deriveKeypair moderateParams

-- | Beep boop
--
-- Note: this function is not thread-safe (since the underlying C function is
-- not thread-safe both in Sodium and in NaCl)! Either make sure
-- there are no concurrent calls or see @Crypto.Init@ in
-- <https://hackage.haskell.org/package/crypto-sodium crypto-sodium> to learn
-- how to make this function thread-safe.
sodiumRekdf ::
  ( ByteArrayAccess pass,
    Sign.Seed ScrubbedBytes Key.!>=! pass
  ) =>
  DerivationSlip ->
  pass ->
  Maybe (PublicKey ByteString, SecretKey ScrubbedBytes)
sodiumRekdf = rederiveKeypair

-- | Generate a new 'SecretKey' together with its 'PublicKey' from passphrase.
--
-- Note: this function is not thread-safe (since the underlying C function is
-- not thread-safe both in Sodium and in NaCl)! Either make sure
-- there are no concurrent calls or see @Crypto.Init@ in
-- <https://hackage.haskell.org/package/crypto-sodium crypto-sodium> to learn
-- how to make this function thread-safe.
--
-- We're using the following underlying libsodium function:

-- $ rg -i crypto_sign_seed_keypair
-- hs-libsodium/libsodium/lib/Libsodium.chs
-- 256:  , crypto_sign_seed_keypair
-- 712:{# fun crypto_sign_seed_keypair { id `pk ::: Ptr CUChar', id `sk ::: Ptr CUChar', id `seed ::: Ptr CUChar' } -> `CInt' #}
--
-- We can use it with @Crypto.Key.Internal@ `derive` because it generates keys
-- that are exacltly `CRYPTO_BOX_SEEDBYTES` (32 at the moment of writing) bytes
-- long, which is exactly the same constant used for `crypto_sign_seed_keypair`.

deriveKeypair ::
  ( Sign.Seed ScrubbedBytes Key.!>=! passwd,
    ByteArray sk,
    ByteArray pk,
    ByteArrayAccess passwd
  ) =>
  Params ->
  passwd ->
  IO (Maybe (PublicKey pk, SecretKey sk, DerivationSlip))
deriveKeypair params pass = do
  seedAndSlipMaybe <- (derive params pass :: IO (Maybe (Sign.Seed ScrubbedBytes, DerivationSlip)))
  withPwhashDeriveKeypair seedAndSlipMaybe

withPwhashDeriveKeypair ::
  ( ByteArray sk,
    ByteArray pk,
    Monad f,
    ByteArrayAccess seed
  ) =>
  Maybe (Sign.Seed seed, c) ->
  f (Maybe (PublicKey pk, SecretKey sk, c))
withPwhashDeriveKeypair Nothing = pure Nothing
withPwhashDeriveKeypair (Just (seed, slip)) = do
  let (pk, sk) = Sign.seededKeypair seed
  return $ Just (pk, sk, slip)

-- | Rederive keypair.
--
-- Note: this function is not thread-safe (since the underlying C function is
-- not thread-safe both in Sodium and in NaCl)! Either make sure
-- there are no concurrent calls or see @Crypto.Init@ in
-- <https://hackage.haskell.org/package/crypto-sodium crypto-sodium> to learn
-- how to make this function thread-safe.
rederiveKeypair ::
  ( ByteArrayAccess pass,
    Sign.Seed ScrubbedBytes Key.!>=! pass
  ) =>
  DerivationSlip ->
  pass ->
  Maybe (PublicKey ByteString, SecretKey ScrubbedBytes)
rederiveKeypair slip pass = do
  withPwhashRederiveKeypair (rederive slip pass :: (Maybe (Sign.Seed ScrubbedBytes)))

withPwhashRederiveKeypair ::
  ( ByteArray sk,
    ByteArray pk,
    ByteArrayAccess seed
  ) =>
  Maybe (Sign.Seed seed) ->
  Maybe (PublicKey pk, SecretKey sk)
withPwhashRederiveKeypair Nothing = Nothing
withPwhashRederiveKeypair (Just seed) = Just $ Sign.seededKeypair seed

-- | PKC-compatible detached signature type
data Signature pk = Signature
  { dacs_pk :: pk,
    dacs_sig :: Sign.Signature ByteString
  }
  deriving (Show, Eq)


-- | Create a detached signature given `PublicKey` and `SecretKey`
sodiumSign ::
  ByteArrayAccess sigdata =>
  (PublicKey ByteString, SecretKey ScrubbedBytes) ->
  sigdata ->
  Signature (PublicKey ByteString)
sodiumSign (pk, sk) msg =
  Signature pk $ Sign.createDetached sk msg

-- | Verify a detached signature
sodiumVerify ::
  ByteArrayAccess sigdata =>
  Signature (PublicKey ByteString) ->
  sigdata ->
  Bool
sodiumVerify (Signature pk sig) msg =
  Sign.verifyDetached sig msg pk

---------- Dummy crypto -----------
data DummyTag = DummySK | DummyPK | DummySlip
  deriving (Show, Eq)

newtype DummySigned b a = DummySigned {dacds_tangledSignature :: (a, (a, b))}
  deriving (Show, Eq)

type DummyTagged = (DummyTag, ByteString)

-- | Dummy cryptographic cunfiguration, used for testing and demonstration
-- purposes.
dummyPkc ::
  PKC
    IO
    ByteString
    DummyTagged
    DummyTagged
    DummyTagged
    (DummySigned ByteString)
    ByteString
dummyPkc =
  PKC
    { dacpkc_kdf = dummyKdf,
      dacpkc_rekdf = dummyRekdf,
      dacpkc_sign = dummySign,
      dacpkc_verify = dummyVerify
    }

dummyKdf :: ByteString -> IO (Maybe (DummyTagged, DummyTagged, DummyTagged))
dummyKdf pass = pure $ Just ((DummyPK, pass), (DummySK, pass), (DummySlip, pass))

dummyRekdf :: DummyTagged -> ByteString -> Maybe (DummyTagged, DummyTagged)
dummyRekdf (DummySlip, x) pass =
  go (x == pass)
  where
    go True = Just ((DummyPK, pass), (DummySK, pass))
    go False = Nothing
dummyRekdf _ _ = error "DummySlip expected in the 1st argument"

dummySign :: (DummyTagged, DummyTagged) -> a -> DummySigned a DummyTagged
dummySign (verificationKey@(DummyPK, _), signingKey@(DummySK, _)) blob =
  DummySigned (verificationKey, (signingKey, blob))
dummySign _ _ = error "Second argument has to be tuple of DummySK and DummyPK"

dummyVerify :: Eq a => DummySigned a DummyTagged -> a -> Bool
dummyVerify (DummySigned ((DummyPK, signedAs), ((DummySK, signedWith), signedWhat))) candidate =
  (signedAs == signedWith) && (candidate == signedWhat)