S2K.hs

-- S2K.hs: OpenPGP (RFC9580) string-to-key conversion
-- Copyright © 2013-2026  Clint Adams
-- This software is released under the terms of the Expat license.
-- (See the LICENSE file).

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}

module Codec.Encryption.OpenPGP.S2K
  ( EncodedSessionKeyError(..)
  , renderEncodedSessionKeyError
  , S2KError(..)
  , renderS2KError
  , decodeOpenPGPEncodedSessionKey
  , string2Key
  , skesk2Key
  , skesk2SessionKey
  ) where

import Codec.Encryption.OpenPGP.BlockCipher (CipherError(..), keySize, withSymmetricCipher)
import Codec.Encryption.OpenPGP.Internal.HOBlockCipher (HOBlockCipher(..))
import Codec.Encryption.OpenPGP.Types
import Data.Bits (shiftL)
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL
import Data.Word (Word8, Word16)
import Crypto.Error (CryptoFailable(..))
import qualified Crypto.Hash as CH
import qualified Crypto.KDF.Argon2 as Argon2
import qualified Data.ByteArray as BA
import Data.Bifunctor (first)

data EncodedSessionKeyError
  = EncodedSessionKeyTooShort
  | EncodedSessionKeyUnsupportedAlgorithm SymmetricAlgorithm
  | EncodedSessionKeyLengthMismatch SymmetricAlgorithm Int Int
  | EncodedSessionKeyChecksumMismatch
  deriving (Eq, Show)

renderEncodedSessionKeyError :: EncodedSessionKeyError -> String
renderEncodedSessionKeyError EncodedSessionKeyTooShort =
  "session key material too short"
renderEncodedSessionKeyError (EncodedSessionKeyUnsupportedAlgorithm sa) =
  "unsupported symmetric algorithm: " ++ show sa
renderEncodedSessionKeyError (EncodedSessionKeyLengthMismatch _ _ _) =
  "session key material length does not match encoded algorithm"
renderEncodedSessionKeyError EncodedSessionKeyChecksumMismatch =
  "session key checksum mismatch"

-- | Errors that can arise during string-to-key derivation.
data S2KError
  = -- | The symmetric algorithm used in the SKESK is not supported.
    S2KUnsupportedAlgorithm CipherError
  | -- | An unsupported or unknown S2K specifier type was encountered.
    S2KUnsupportedSpecifier Word8
  | -- | An unsupported SKESK shape (e.g. non-zero ESK).
    S2KUnsupportedSKESKShape String
  | -- | A required hash algorithm is not supported for S2K.
    S2KUnsupportedHashAlgorithm HashAlgorithm
  | -- | The Argon2 S2K parameters are invalid.
    S2KArgon2ParamError String
  | -- | The Argon2 KDF itself failed.
    S2KArgon2Failed String
  | -- | Decrypting an embedded encrypted session key failed.
    S2KEncryptedSessionKeyCipherError CipherError
  | -- | Embedded encrypted session key material was malformed.
    S2KEncryptedSessionKeyDecodeError EncodedSessionKeyError
  deriving (Eq, Show)

renderS2KError :: S2KError -> String
renderS2KError (S2KUnsupportedAlgorithm ce) =
  "S2K: " ++ renderCipherError' ce
  where
    renderCipherError' (UnsupportedAlgorithm sa) = "unsupported symmetric algorithm: " ++ show sa
    renderCipherError' (CipherInitFailed sa msg) = "cipher init failed for " ++ show sa ++ ": " ++ msg
    renderCipherError' (CipherOperationFailed msg) = "cipher operation failed: " ++ msg
renderS2KError (S2KUnsupportedSpecifier t) =
  "S2K: unsupported S2K type " ++ show t
renderS2KError (S2KUnsupportedSKESKShape msg) =
  "S2K: unsupported SKESK shape: " ++ msg
renderS2KError (S2KUnsupportedHashAlgorithm ha) =
  "S2K: unsupported hash algorithm for S2K: " ++ show ha
renderS2KError (S2KArgon2ParamError msg) =
  "S2K: Argon2 parameter error: " ++ msg
renderS2KError (S2KArgon2Failed msg) =
  "S2K: Argon2 KDF failed: " ++ msg
renderS2KError (S2KEncryptedSessionKeyCipherError ce) =
  "S2K: encrypted session key decrypt failed: " ++ renderCipherError' ce
  where
    renderCipherError' (UnsupportedAlgorithm sa) = "unsupported symmetric algorithm: " ++ show sa
    renderCipherError' (CipherInitFailed sa msg) = "cipher init failed for " ++ show sa ++ ": " ++ msg
    renderCipherError' (CipherOperationFailed msg) = "cipher operation failed: " ++ msg
renderS2KError (S2KEncryptedSessionKeyDecodeError err) =
  "S2K: encrypted session key decode failed: " ++ renderEncodedSessionKeyError err

string2Key :: S2K -> Int -> BL.ByteString -> Either S2KError B.ByteString
string2Key (Simple ha) ksz bs =
  B.take (fromIntegral ksz) <$> hashpp ha ksz bs
string2Key (Salted ha salt) ksz bs =
  string2Key (Simple ha) ksz (BL.append (BL.fromStrict (unSalt8 salt)) bs)
string2Key (IteratedSalted ha salt cnt) ksz bs =
  string2Key
    (Simple ha)
    ksz
    (BL.take (fromIntegral cnt) . BL.cycle $
     BL.append (BL.fromStrict (unSalt8 salt)) bs)
string2Key (Argon2 salt t p encodedM) ksz pass =
  argon2String2Key salt t p encodedM ksz pass
string2Key (OtherS2K t _) _ _ =
  Left (S2KUnsupportedSpecifier t)

skesk2Key :: SKESK 'SKESKV4 -> BL.ByteString -> Either S2KError B.ByteString
skesk2Key skesk pass = snd <$> skesk2SessionKey skesk pass

skesk2SessionKey :: SKESK 'SKESKV4 -> BL.ByteString -> Either S2KError (SymmetricAlgorithm, B.ByteString)
skesk2SessionKey (SKESK4Packet sa s2k Nothing) pass = do
  keyLen <- first S2KUnsupportedAlgorithm (keySize sa)
  sessionKey <- string2Key s2k keyLen pass
  pure (sa, sessionKey)
skesk2SessionKey (SKESK4Packet sa s2k (Just esk)) pass = do
  keyLen <- first S2KUnsupportedAlgorithm (keySize sa)
  kek <- string2Key s2k keyLen pass
  decrypted <-
    first S2KEncryptedSessionKeyCipherError $
    withSymmetricCipher sa kek
      (\cipher ->
         paddedCfbDecrypt
           cipher
           (B.replicate (blockSize cipher) 0)
           (BL.toStrict esk))
  first S2KEncryptedSessionKeyDecodeError (decodeSKESK4EncryptedSessionKey decrypted)

decodeOpenPGPEncodedSessionKey ::
     B.ByteString -> Either EncodedSessionKeyError (SymmetricAlgorithm, B.ByteString)
decodeOpenPGPEncodedSessionKey encodedSessionKey = do
  if B.length encodedSessionKey < 3
    then Left EncodedSessionKeyTooShort
    else Right ()
  let symalgo = toFVal (B.head encodedSessionKey)
      rest = B.tail encodedSessionKey
  keyLen <- encodedSessionKeyKeyLength symalgo
  if B.length rest < keyLen + 2
    then Left (EncodedSessionKeyLengthMismatch symalgo keyLen (B.length rest))
    else Right ()
  let (sessionKey, restAfterKey) = B.splitAt keyLen rest
      checksumBytes = B.take 2 restAfterKey
      actualChecksum = checksum16 sessionKey
      expectedChecksum =
        fromIntegral (B.index checksumBytes 0) `shiftL` 8 +
        fromIntegral (B.index checksumBytes 1)
  if actualChecksum /= expectedChecksum
    then Left EncodedSessionKeyChecksumMismatch
    else Right (symalgo, sessionKey)

decodeSKESK4EncryptedSessionKey ::
     B.ByteString -> Either EncodedSessionKeyError (SymmetricAlgorithm, B.ByteString)
decodeSKESK4EncryptedSessionKey encodedSessionKey = do
  if B.length encodedSessionKey < 1
    then Left EncodedSessionKeyTooShort
    else Right ()
  let symalgo = toFVal (B.head encodedSessionKey)
      sessionKey = B.tail encodedSessionKey
  keyLen <- encodedSessionKeyKeyLength symalgo
  if B.length sessionKey /= keyLen
    then Left (EncodedSessionKeyLengthMismatch symalgo keyLen (B.length sessionKey))
    else Right (symalgo, sessionKey)

encodedSessionKeyKeyLength ::
     SymmetricAlgorithm -> Either EncodedSessionKeyError Int
encodedSessionKeyKeyLength symalgo =
  first renderKeySizeError (keySize symalgo)
  where
    renderKeySizeError :: CipherError -> EncodedSessionKeyError
    renderKeySizeError (UnsupportedAlgorithm sa) =
      EncodedSessionKeyUnsupportedAlgorithm sa
    renderKeySizeError (CipherInitFailed sa _) =
      EncodedSessionKeyUnsupportedAlgorithm sa
    renderKeySizeError (CipherOperationFailed _) =
      EncodedSessionKeyUnsupportedAlgorithm symalgo

checksum16 :: B.ByteString -> Word16
checksum16 =
  fromIntegral .
  B.foldl' (\acc octet -> (acc + fromIntegral octet) `mod` (65536 :: Integer)) 0

argon2String2Key :: Salt16 -> Word8 -> Word8 -> Word8 -> Int -> BL.ByteString -> Either S2KError B.ByteString
argon2String2Key salt t p encodedM keyLen pass
  | t == 0 = Left (S2KArgon2ParamError "Argon2 S2K pass count must be non-zero")
  | p == 0 = Left (S2KArgon2ParamError "Argon2 S2K parallelism must be non-zero")
  | encodedM > 31 = Left (S2KArgon2ParamError "Argon2 S2K encoded_m must be <= 31")
  | encodedM < minEncodedM = Left (S2KArgon2ParamError "Argon2 S2K encoded_m is too small for parallelism")
  | otherwise =
      case Argon2.hash opts (BL.toStrict pass) (unSalt16 salt) keyLen of
        CryptoPassed k -> Right k
        CryptoFailed e -> Left (S2KArgon2Failed (show e))
  where
    opts =
      Argon2.defaultOptions
        { Argon2.iterations = fromIntegral t
        , Argon2.memory = fromIntegral (1 `shiftL` fromIntegral encodedM :: Int)
        , Argon2.parallelism = fromIntegral p
        , Argon2.variant = Argon2.Argon2id
        , Argon2.version = Argon2.Version13
        }
    minEncodedM = fromIntegral (3 + ceilLog2 (fromIntegral p :: Int))

ceilLog2 :: Int -> Int
ceilLog2 n
  | n <= 1 = 0
  | otherwise = go 0 1
  where
    go e v
      | v >= n = e
      | otherwise = go (e + 1) (v * 2)

hashpp :: HashAlgorithm -> Int -> BL.ByteString -> Either S2KError B.ByteString
hashpp ha keysize pp =
  go 0 B.empty
  where
    go ctr acc
      | B.length acc >= keysize = Right acc
      | otherwise = do
          digest <- hf ha (nulpad ctr `BL.append` pp)
          go (ctr + 1) (acc `B.append` digest)
    nulpad = BL.pack . flip replicate 0
    hf :: HashAlgorithm -> BL.ByteString -> Either S2KError B.ByteString
    hf DeprecatedMD5 bs = Right (BA.convert (CH.hashlazy bs :: CH.Digest CH.MD5))
    hf SHA1 bs = Right (BA.convert (CH.hashlazy bs :: CH.Digest CH.SHA1))
    hf SHA224 bs = Right (BA.convert (CH.hashlazy bs :: CH.Digest CH.SHA224))
    hf SHA256 bs = Right (BA.convert (CH.hashlazy bs :: CH.Digest CH.SHA256))
    hf SHA384 bs = Right (BA.convert (CH.hashlazy bs :: CH.Digest CH.SHA384))
    hf SHA512 bs = Right (BA.convert (CH.hashlazy bs :: CH.Digest CH.SHA512))
    hf SHA3_256 bs = Right (BA.convert (CH.hashlazy bs :: CH.Digest CH.SHA3_256))
    hf SHA3_512 bs = Right (BA.convert (CH.hashlazy bs :: CH.Digest CH.SHA3_512))
    hf (OtherHA ha') _ = Left (S2KUnsupportedHashAlgorithm (OtherHA ha'))