content: {-# 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 Irc.Commands (AuthenticatePayload (AuthenticatePayload))
  import OpenSSL.EVP.Digest (Digest, digestBS, hmacBS, getDigestByName)
  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
  