{-# Language BlockArguments #-}
{-# Language ImportQualifiedPost #-}
{-# Language LambdaCase #-}
{-# Language OverloadedStrings #-}
{-# Language RecordWildCards #-}
{-# Language ViewPatterns #-}
module Client.Authentication.Scram (
  -- * Transaction state types
  Phase1,
  Phase2,
  -- * Transaction step functions
  initiateScram,
  addServerFirst,
  addServerFinal,
  -- * Digests
  ScramDigest(..),
  mechanismName,
  ) where

import Control.Monad (guard)
import Data.Bits (xor)
import Data.ByteString (ByteString)
import Data.ByteString qualified as B
import Data.ByteString.Base64 qualified as B64
import Data.ByteString.Char8 qualified as B8
import Data.List (foldl1')
import Data.Text (Text)
import OpenSSL.EVP.Digest ( Digest, digestBS, hmacBS, getDigestByName)
import Irc.Commands (AuthenticatePayload (AuthenticatePayload))
import System.IO.Unsafe (unsafePerformIO)

data ScramDigest
  = ScramDigestSha1
  | ScramDigestSha2_256
  | ScramDigestSha2_512
  deriving Show

mechanismName :: ScramDigest -> Text
mechanismName digest =
  case digest of
    ScramDigestSha1     -> "SCRAM-SHA-1"
    ScramDigestSha2_256 -> "SCRAM-SHA-256"
    ScramDigestSha2_512 -> "SCRAM-SHA-512"

-- | SCRAM state waiting for server-first-message
data Phase1 = Phase1
  { phase1Digest          :: ScramDigest -- ^ underlying cryptographic hash function
  , phase1Password        :: ByteString -- ^ password
  , phase1CbindInput      :: ByteString -- ^ cbind-input
  , phase1Nonce           :: ByteString -- ^ c-nonce
  , phase1ClientFirstBare :: ByteString -- ^ client-first-bare
  }

-- | Construct client-first-message and extra parameters
-- needed for 'addServerFirst'.
initiateScram ::
  ScramDigest ->
  ByteString {- ^ authentication ID -} ->
  ByteString {- ^ authorization ID  -} ->
  ByteString {- ^ password          -} ->
  ByteString {- ^ nonce             -} ->
  (AuthenticatePayload, Phase1)
initiateScram digest user authzid pass nonce =
  (AuthenticatePayload clientFirstMessage, Phase1
    { phase1Digest = digest
    , phase1Password = pass
    , phase1CbindInput = B64.encode gs2Header
    , phase1Nonce = nonce
    , phase1ClientFirstBare = clientFirstMessageBare
    })
  where
    clientFirstMessage = gs2Header <> clientFirstMessageBare
    gs2Header = "n," <> encodeUsername authzid <> ","
    clientFirstMessageBare = "n=" <> encodeUsername user <> ",r=" <> nonce

-- | SCRAM state waiting for server-final-message
newtype Phase2 = Phase2
  { phase2ServerSignature :: ByteString -- ^ base64 encoded expected value
  }

-- | Add server-first-message to current SCRAM transaction,
-- compute client-final-message and next state for 'addServerFinal'.
addServerFirst ::
  Phase1     {- ^ output of 'initiateScram' -} ->
  ByteString {- ^ server first message -} ->
  Maybe (AuthenticatePayload, Phase2)
addServerFirst Phase1{..} serverFirstMessage =

  do -- Parse server-first-message
     ("r", nonce) :
       ("s", B64.decode -> Right salt) :
       ("i", B8.readInt -> Just (iterations, "")) :
       _extensions
       <- Just (parseMessage serverFirstMessage)

     -- validate nonce given by server includes ours and isn't empty
     guard (B.isPrefixOf phase1Nonce nonce && phase1Nonce /= nonce)

     let clientFinalWithoutProof = "c=" <> phase1CbindInput <> ",r=" <> nonce

     let authMessage =
           phase1ClientFirstBare <> "," <>
           serverFirstMessage <> "," <>
           clientFinalWithoutProof

     let (clientProof, serverSignature) =
           crypto phase1Digest phase1Password salt iterations authMessage

     let proof = "p=" <> B64.encode clientProof
     let clientFinalMessage = clientFinalWithoutProof <> "," <> proof

     let phase2 = Phase2 { phase2ServerSignature = B64.encode serverSignature }
     Just (AuthenticatePayload clientFinalMessage, phase2)

-- | Add server-final-message to transaction and compute validatity of
-- the whole transaction.
addServerFinal ::
  Phase2     {- ^ output of 'addServerFirst' -} ->
  ByteString {- ^ server-final-message   -} ->
  Bool       {- ^ transaction succeeded? -}
addServerFinal Phase2{..} serverFinalMessage =
  case parseMessage serverFinalMessage of
    ("v", sig) : _extensions -> sig == phase2ServerSignature
    _ -> False

-- | Big endian encoding of a 32-bit number 1.
int1 :: ByteString
int1 = B.pack [0,0,0,1]

xorBS :: ByteString -> ByteString -> ByteString
xorBS x y = B.pack (B.zipWith xor x y)

-- | Iterated, password-based, key-derivation function.
hi ::
  Digest     {- ^ underlying cryptographic hash function -} ->
  ByteString {- ^ secret -} ->
  ByteString {- ^ salt -} ->
  Int        {- ^ iterations -} ->
  ByteString {- ^ salted, iterated hash of secret -}
hi digest str salt n = foldl1' xorBS (take n us)
  where
    u1 = hmacBS digest str (salt <> int1)
    us = iterate (hmacBS digest str) u1

-- | Break up a SCRAM message into its underlying key-value association list.
parseMessage :: ByteString -> [(ByteString, ByteString)]
parseMessage msg =
  [case B8.break ('='==) entry of
     (key, value) -> (key, B.drop 1 value)
  | entry <- B8.split ',' msg]

-- | Tranform all the SCRAM parameters into a @ClientProof@
-- and @ServerSignature@.
crypto ::
  ScramDigest {- ^ digest       -} ->
  ByteString  {- ^ password     -} ->
  ByteString  {- ^ salt         -} ->
  Int         {- ^ iterations   -} ->
  ByteString  {- ^ auth message -} ->
  (ByteString, ByteString) {- ^ client-proof, server-signature -}
crypto digest password salt iterations authMessage =
  (clientProof, serverSignature)
  where
    saltedPassword  = hi       d password salt iterations
    clientKey       = hmacBS   d saltedPassword "Client Key"
    storedKey       = digestBS d clientKey
    clientSignature = hmacBS   d storedKey authMessage
    clientProof     = xorBS clientKey clientSignature
    serverKey       = hmacBS   d saltedPassword "Server Key"
    serverSignature = hmacBS   d serverKey authMessage
    digestName =
      case digest of
        ScramDigestSha1     -> "SHA1"
        ScramDigestSha2_256 -> "SHA256"
        ScramDigestSha2_512 -> "SHA512"
    Just d = unsafePerformIO (getDigestByName digestName)

-- | Encode usersnames so they fit in the comma/equals delimited
-- SCRAM message format.
encodeUsername :: ByteString -> ByteString
encodeUsername = B8.concatMap \case
    ',' -> "=2C"
    '=' -> "=3D"
    x   -> B8.singleton x