-- Properties.hs: hOpenPGP test suite
-- Copyright © 2012-2026  Clint Adams
-- This software is released under the terms of the Expat license.
-- (See the LICENSE file).
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE OverloadedStrings #-}

module Tests.Properties (propertiesTests) where

import Control.Exception (SomeException, try)
import Control.Lens (preview)
import Data.Binary (get, put)
import Data.Binary.Put (runPut)
import qualified Data.ByteString.Lazy as BL
import qualified Data.Conduit as DC
import qualified Data.Conduit.List as CL
import Data.Maybe (isNothing)
import Data.Word (Word8)
import Test.Tasty (TestTree, localOption, testGroup)
import qualified Test.Tasty.QuickCheck as QC

import Codec.Encryption.OpenPGP.Encrypt
    ( canonicalizePKESKRecipientId
    )
import Codec.Encryption.OpenPGP.KeyringParser
    ( parseTKsWithWireRep
    )
import Codec.Encryption.OpenPGP.Policy
    ( OpenPGPPolicy (..)
    , OpenPGPRFC (RFC4880)
    , defaultPolicy
    , policyForRFC
    )
import Codec.Encryption.OpenPGP.SecretKey
    ( decryptPrivateKey
    , encryptPrivateKey
    )
import Codec.Encryption.OpenPGP.Serialize
    ( parsePktsEither
    , parsePktsWithWireRep
    )
import Codec.Encryption.OpenPGP.Types
import Tests.Common
    ( collectSecretKeyInfos
    , conduitDecryptWithPKESKContext
    , loadSEIPDv2FixtureWithV4Secret
    , loadV4EncryptedSecretKeyFixtureForProperty
    , loadV6UnencryptedSecretKeyFixtureForProperty
    , prependUnusableLatestPKESK
    , readFixtureLazy
    , reorderPrecedingPKESKs
    , reverseIf
    , runGet
    , selectRecipientKeyInfo
    )

propertiesTests :: TestTree
propertiesTests = testGroup "Properties" [qcProps]

qcProps :: TestTree
qcProps =
    testGroup
        "(checked by QuickCheck)"
        [ QC.testProperty "PKESKv3 packet serialization-deserialization" $ \pkesk ->
            Right (pkesk :: PKESK 'PKESKV3)
                == runGet get (runPut (put pkesk))
        , QC.testProperty "PKESKv6 packet serialization-deserialization" $ \pkesk ->
            Right (pkesk :: PKESK 'PKESKV6)
                == runGet get (runPut (put pkesk))
        , QC.testProperty "Signature packet serialization-deserialization" $ \sig ->
            ( isNothing
                (preview _SigVOther (_signaturePayload (sig :: Signature)))
            )
                QC.==> Right (sig :: Signature) == runGet get (runPut (put sig))
        , QC.testProperty "UserId packet serialization-deserialization" $ \uid ->
            Right (uid :: UserId) == runGet get (runPut (put uid))
        , QC.testProperty
            "decryptPrivateKey (encryptPrivateKey sk pw) pw equivalence"
            $ \passphraseNE ->
                QC.ioProperty $ do
                    fixture <- loadV6UnencryptedSecretKeyFixtureForProperty
                    case fixture of
                        Left err -> pure (QC.counterexample err False)
                        Right (pkp, ska, expectedSKey) -> do
                            let passphraseChars = (QC.getNonEmpty passphraseNE :: String)
                                passphrase = BL.pack (map (fromIntegral . fromEnum) passphraseChars)
                            encryptedResult <-
                                encryptPrivateKey defaultPolicy pkp ska passphrase
                            pure $
                                case encryptedResult of
                                    Left err ->
                                        QC.counterexample ("encryptPrivateKey failed: " ++ err) False
                                    Right encryptedSKA ->
                                        case decryptPrivateKey (pkp, encryptedSKA) passphrase of
                                            Left err ->
                                                QC.counterexample ("decryptPrivateKey failed: " ++ err) False
                                            Right (SUUnencrypted skey _) ->
                                                QC.counterexample
                                                    "secret key material changed across encrypt/decrypt roundtrip"
                                                    (skey == expectedSKey)
                                            Right other ->
                                                QC.counterexample
                                                    ( "expected SUUnencrypted after decrypting encrypted key, got: "
                                                        ++ show other
                                                    )
                                                    False
        , localOption
            (QC.QuickCheckTests 5)
            ( QC.testProperty
                "decryptPrivateKey (encryptPrivateKey v4sk pw) pw equivalence"
                $ QC.ioProperty
                $ do
                    fixture <- loadV4EncryptedSecretKeyFixtureForProperty
                    case fixture of
                        Left err -> pure (QC.counterexample err False)
                        Right (pkp, ska, expectedSKey, passphrase) -> do
                            let legacyOverridePolicy :: OpenPGPPolicy
                                legacyOverridePolicy =
                                    (policyForRFC RFC4880)
                                        { policySecretKeyProtection =
                                            policySecretKeyProtection defaultPolicy
                                        }
                            encryptedResult <-
                                encryptPrivateKey legacyOverridePolicy pkp ska passphrase
                            pure $
                                case encryptedResult of
                                    Left err ->
                                        QC.counterexample
                                            ("encryptPrivateKey failed under legacy override policy: " ++ err)
                                            False
                                    Right encryptedSKA ->
                                        case decryptPrivateKey (pkp, encryptedSKA) passphrase of
                                            Left err ->
                                                QC.counterexample ("decryptPrivateKey failed: " ++ err) False
                                            Right (SUUnencrypted skey _) ->
                                                QC.counterexample
                                                    "v4 secret key material changed across encrypt/decrypt roundtrip"
                                                    (skey == expectedSKey)
                                            Right other ->
                                                QC.counterexample
                                                    ( "expected SUUnencrypted after decrypting v4 encrypted key, got: "
                                                        ++ show other
                                                    )
                                                    False
            )
        , QC.testProperty
            "canonicalizePKESKRecipientId idempotence on valid v4/v6 recipient ids"
            propertyCanonicalizePKESKRecipientIdIdempotent
        , localOption
            (QC.QuickCheckTests 10)
            ( QC.testProperty
                "canonicalizeTKStructuredWithWireRep is stable across packet/sig reordering"
                propertyCanonicalizeTKStructuredStableAcrossReordering
            )
        , localOption
            (QC.QuickCheckTests 1)
            ( QC.testProperty
                "recipient selection remains decryptable with extra unusable PKESK candidates"
                propertyRecipientSelectionMonotonicWithUnusableCandidates
            )
        , QC.testProperty
            "parsePktsEither does not accept truncated packet streams as intact packets"
            propertyParsePktsEitherRejectsTruncatedPacketStream
        ]

propertyCanonicalizePKESKRecipientIdIdempotent
    :: Bool -> Bool -> [Word8] -> QC.Property
propertyCanonicalizePKESKRecipientIdIdempotent useV6 prefixed seedBytes =
    case canonicalizePKESKRecipientId payload of
        Left err ->
            QC.counterexample
                ("canonicalizePKESKRecipientId unexpectedly failed: " ++ show err)
                False
        Right canonical ->
            QC.counterexample
                "canonicalizePKESKRecipientId should be idempotent"
                (canonicalizePKESKRecipientId canonical == Right canonical)
  where
    targetLen = if useV6 then 32 else 20
    versionOctet = if useV6 then 0x06 else 0x04
    ridBody = BL.pack (take targetLen (seedBytes ++ repeat 0x00))
    rid
        | prefixed = BL.cons versionOctet ridBody
        | otherwise = ridBody
    payload = PKESKPayloadV6Packet (PKESKPayloadV6 rid RSA "esk")

propertyCanonicalizeTKStructuredStableAcrossReordering
    :: QC.NonNegative Int
    -> Bool
    -> Bool
    -> Bool
    -> Bool
    -> Bool
    -> Bool
    -> Bool
    -> QC.Property
propertyCanonicalizeTKStructuredStableAcrossReordering
    (QC.NonNegative indexSeed)
    reverseDirect
    reverseUIDs
    reverseUIDSigs
    reverseUATs
    reverseUATSigs
    reverseSubs
    reverseSubSigs =
        QC.ioProperty $ do
            lbs <- readFixtureLazy "pubring.gpg"
            let src = wireRepRef lbs
                parsed = parseTKsWithWireRep True (parsePktsWithWireRep src lbs)
            if null parsed
                then
                    pure
                        ( QC.counterexample
                            "pubring.gpg parsed to no TKWithWireRep values"
                            False
                        )
                else do
                    let tk = parsed !! (indexSeed `mod` length parsed)
                    pure $
                        case toStructuredTKWithWireRep tk of
                            Left err ->
                                QC.counterexample
                                    ("toStructuredTKWithWireRep failed: " ++ err)
                                    False
                            Right structured ->
                                let shuffled =
                                        structured
                                            { _tkStructuredDirectSignatures =
                                                reverseIf
                                                    reverseDirect
                                                    (_tkStructuredDirectSignatures structured)
                                            , _tkStructuredUIDs =
                                                reverseIf
                                                    reverseUIDs
                                                    ( map
                                                        ( \uid ->
                                                            uid
                                                                { _uidWithWireRefsSignatures =
                                                                    reverseIf reverseUIDSigs (_uidWithWireRefsSignatures uid)
                                                                }
                                                        )
                                                        (_tkStructuredUIDs structured)
                                                    )
                                            , _tkStructuredUAts =
                                                reverseIf
                                                    reverseUATs
                                                    ( map
                                                        ( \uat ->
                                                            uat
                                                                { _uatWithWireRefsSignatures =
                                                                    reverseIf reverseUATSigs (_uatWithWireRefsSignatures uat)
                                                                }
                                                        )
                                                        (_tkStructuredUAts structured)
                                                    )
                                            , _tkStructuredSubkeys =
                                                reverseIf
                                                    reverseSubs
                                                    ( map
                                                        ( \sub ->
                                                            sub
                                                                { _subkeyWithWireRefsSignatures =
                                                                    reverseIf reverseSubSigs (_subkeyWithWireRefsSignatures sub)
                                                                }
                                                        )
                                                        (_tkStructuredSubkeys structured)
                                                    )
                                            }
                                 in case ( canonicalizeTKStructuredWithWireRep structured
                                         , canonicalizeTKStructuredWithWireRep shuffled
                                         ) of
                                        (Right canonicalBase, Right canonicalShuffled) ->
                                            QC.counterexample
                                                "canonicalization should be stable under packet/sig reordering"
                                                (canonicalBase == canonicalShuffled)
                                        (Left err, _) ->
                                            QC.counterexample
                                                ( "canonicalizeTKStructuredWithWireRep failed on base: "
                                                    ++ show err
                                                )
                                                False
                                        (_, Left err) ->
                                            QC.counterexample
                                                ( "canonicalizeTKStructuredWithWireRep failed on shuffled: "
                                                    ++ show err
                                                )
                                                False

propertyRecipientSelectionMonotonicWithUnusableCandidates
    :: QC.NonNegative Int -> Bool -> QC.Property
propertyRecipientSelectionMonotonicWithUnusableCandidates (QC.NonNegative extraBogus) reorderPKESKs =
    QC.ioProperty $ do
        let applyBogus =
                foldr
                    (.)
                    id
                    (replicate (extraBogus `mod` 5) prependUnusableLatestPKESK)
            transformPackets
                | reorderPKESKs = applyBogus . reorderPrecedingPKESKs
                | otherwise = applyBogus
        result <-
            ( try $ do
                (messagePacketsRaw, encryptedSecretPackets, passphrase) <-
                    loadSEIPDv2FixtureWithV4Secret "seipdv2-three-recipients.pgp.aa"
                keyInfos <-
                    collectSecretKeyInfos encryptedSecretPackets passphrase
                let messagePackets = transformPackets messagePacketsRaw
                    passphraseCallback _ = pure BL.empty
                    keyContextCallback pkt = pure (selectRecipientKeyInfo pkt keyInfos)
                decrypted <-
                    DC.runConduitRes $
                        CL.sourceList messagePackets
                            DC..| conduitDecryptWithPKESKContext
                                keyContextCallback
                                passphraseCallback
                            DC..| CL.consume
                pure
                    ( any
                        (not . BL.null)
                        [payload | LiteralDataPkt _ _ _ payload <- decrypted]
                    )
            )
                :: IO (Either SomeException Bool)
        case result of
            Left e ->
                pure $
                    QC.counterexample
                        ( "recipient selection should remain decryptable despite added unusable candidates: "
                            ++ show e
                        )
                        False
            Right didDecrypt ->
                pure $
                    QC.counterexample
                        "recipient selection should still yield a non-empty decrypted literal payload"
                        didDecrypt

propertyParsePktsEitherRejectsTruncatedPacketStream
    :: PKESK 'PKESKV6 -> QC.Positive Int -> QC.Property
propertyParsePktsEitherRejectsTruncatedPacketStream pkesk (QC.Positive cutSeed) =
    case parsePktsEither truncated of
        Left _ -> QC.property True
        Right parsed ->
            QC.counterexample
                ( "parsePktsEither unexpectedly treated truncated packet stream as original packet: "
                    ++ show parsed
                )
                (parsed /= [toPkt pkesk])
  where
    encoded = runPut (put (toPkt pkesk))
    cut =
        fromIntegral
            (1 + (cutSeed `mod` fromIntegral (BL.length encoded)))
    truncated = BL.take (BL.length encoded - cut) encoded
