-- Serialization.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 OverloadedStrings #-}

module Tests.Serialization (serializationTests) where

import Control.Applicative ((<|>))
import Control.Lens ((^.))
import Control.Monad (forM_)
import Crypto.Number.Serialize (os2ip)
import qualified Crypto.PubKey.ECC.ECDSA as ECDSA
import qualified Crypto.PubKey.ECC.Types as ECCT
import Data.Binary (Get, get, put)
import Data.Binary.Put
    ( putByteString
    , putWord16be
    , putWord32be
    , putWord8
    , runPut
    )
import Data.Bits (xor)
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL
import qualified Data.Conduit as DC
import qualified Data.Conduit.Binary as CB
import qualified Data.Conduit.List as CL
import Data.Conduit.Serialization.Binary (conduitGet)
import Data.List.NonEmpty (NonEmpty (..))
import qualified Data.List.NonEmpty as NE
import Data.Maybe (listToMaybe)
import qualified Data.Set as Set
import Data.Word (Word8)
import Test.Tasty (TestTree, testGroup)
import Test.Tasty.HUnit
    ( Assertion
    , assertBool
    , assertEqual
    , assertFailure
    , testCase
    )

import Codec.Encryption.OpenPGP.Fingerprint
    ( eightOctetKeyID
    , fingerprint
    )
import Codec.Encryption.OpenPGP.Internal (emptyPSC, point2MBS)
import Codec.Encryption.OpenPGP.KeyringParser
    ( parseTKsWithWireRep
    )
import Codec.Encryption.OpenPGP.Policy
    ( defaultVerificationPolicy
    , signatureV6SaltSizeForHashAlgorithm
    )
import Codec.Encryption.OpenPGP.Serialize
    ( parsePkts
    , parsePktsWithWireRep
    )
import Codec.Encryption.OpenPGP.Signatures
    ( VerificationError (..)
    , verifySigWith
    )
import Codec.Encryption.OpenPGP.Types
import qualified Codec.Encryption.OpenPGP.Types.Internal.Base as PKA
import Tests.Common
    ( armorPayload
    , loadArmor
    , readFixturePayload
    , runGet
    )

serializationTests :: TestTree
serializationTests =
    testGroup
        "Serialization"
        [ testGroup
            "Serialization group"
            [ testCase
                "000001-006.public_key"
                (testSerialization "000001-006.public_key")
            , testCase
                "issuer-fingerprint-rejects-unknown-version"
                testIssuerFingerprintRejectsUnknownVersion
            , testCase
                "v6-ed25519-public-key-serializes-fixed-length"
                testV6Ed25519PublicKeySerializesFixedLength
            , testCase
                "v4-ed25519-public-key-parses-native-fixed-length"
                testV4Ed25519PublicKeyParsesNativeFixedLength
            , testCase
                "v4-x25519-public-subkey-parses-native-fixed-length"
                testV4X25519PublicSubkeyParsesNativeFixedLength
            , testCase
                "v6-signature-issuer-fingerprint-version-mismatch"
                testV6SignatureIssuerFingerprintVersionMismatch
            , testCase
                "000002-013.user_id"
                (testSerialization "000002-013.user_id")
            , testCase "000003-002.sig" (testSerialization "000003-002.sig")
            , testCase
                "000004-012.ring_trust"
                (testSerialization "000004-012.ring_trust")
            , testCase "000005-002.sig" (testSerialization "000005-002.sig")
            , testCase
                "000006-012.ring_trust"
                (testSerialization "000006-012.ring_trust")
            , testCase "000007-002.sig" (testSerialization "000007-002.sig")
            , testCase
                "000008-012.ring_trust"
                (testSerialization "000008-012.ring_trust")
            , testCase "000009-002.sig" (testSerialization "000009-002.sig")
            , testCase
                "000010-012.ring_trust"
                (testSerialization "000010-012.ring_trust")
            , testCase "000011-002.sig" (testSerialization "000011-002.sig")
            , testCase
                "000012-012.ring_trust"
                (testSerialization "000012-012.ring_trust")
            , testCase
                "000013-014.public_subkey"
                (testSerialization "000013-014.public_subkey")
            , testCase "000014-002.sig" (testSerialization "000014-002.sig")
            , testCase
                "000015-012.ring_trust"
                (testSerialization "000015-012.ring_trust")
            , testCase
                "000016-006.public_key"
                (testSerialization "000016-006.public_key")
            , testCase "000017-002.sig" (testSerialization "000017-002.sig")
            , testCase
                "000018-012.ring_trust"
                (testSerialization "000018-012.ring_trust")
            , testCase
                "000019-013.user_id"
                (testSerialization "000019-013.user_id")
            , testCase "000020-002.sig" (testSerialization "000020-002.sig")
            , testCase
                "000021-012.ring_trust"
                (testSerialization "000021-012.ring_trust")
            , testCase "000022-002.sig" (testSerialization "000022-002.sig")
            , testCase
                "000023-012.ring_trust"
                (testSerialization "000023-012.ring_trust")
            , testCase
                "000024-014.public_subkey"
                (testSerialization "000024-014.public_subkey")
            , testCase "000025-002.sig" (testSerialization "000025-002.sig")
            , testCase
                "000026-012.ring_trust"
                (testSerialization "000026-012.ring_trust")
            , testCase
                "000027-006.public_key"
                (testSerialization "000027-006.public_key")
            , testCase "000028-002.sig" (testSerialization "000028-002.sig")
            , testCase
                "000029-012.ring_trust"
                (testSerialization "000029-012.ring_trust")
            , testCase
                "000030-013.user_id"
                (testSerialization "000030-013.user_id")
            , testCase "000031-002.sig" (testSerialization "000031-002.sig")
            , testCase
                "000032-012.ring_trust"
                (testSerialization "000032-012.ring_trust")
            , testCase "000033-002.sig" (testSerialization "000033-002.sig")
            , testCase
                "000034-012.ring_trust"
                (testSerialization "000034-012.ring_trust")
            , testCase
                "000035-006.public_key"
                (testSerialization "000035-006.public_key")
            , testCase
                "000036-013.user_id"
                (testSerialization "000036-013.user_id")
            , testCase "000037-002.sig" (testSerialization "000037-002.sig")
            , testCase
                "000038-012.ring_trust"
                (testSerialization "000038-012.ring_trust")
            , testCase "000039-002.sig" (testSerialization "000039-002.sig")
            , testCase
                "000040-012.ring_trust"
                (testSerialization "000040-012.ring_trust")
            , testCase
                "000041-017.attribute"
                (testSerialization "000041-017.attribute")
            , testCase "000042-002.sig" (testSerialization "000042-002.sig")
            , testCase
                "000043-012.ring_trust"
                (testSerialization "000043-012.ring_trust")
            , testCase
                "000044-014.public_subkey"
                (testSerialization "000044-014.public_subkey")
            , testCase "000045-002.sig" (testSerialization "000045-002.sig")
            , testCase
                "000046-012.ring_trust"
                (testSerialization "000046-012.ring_trust")
            , testCase
                "000047-005.secret_key"
                (testSerialization "000047-005.secret_key")
            , testCase
                "000048-013.user_id"
                (testSerialization "000048-013.user_id")
            , testCase "000049-002.sig" (testSerialization "000049-002.sig")
            , testCase
                "000050-012.ring_trust"
                (testSerialization "000050-012.ring_trust")
            , testCase
                "000051-007.secret_subkey"
                (testSerialization "000051-007.secret_subkey")
            , testCase "000052-002.sig" (testSerialization "000052-002.sig")
            , testCase
                "000053-012.ring_trust"
                (testSerialization "000053-012.ring_trust")
            , testCase
                "000054-005.secret_key"
                (testSerialization "000054-005.secret_key")
            , testCase "000055-002.sig" (testSerialization "000055-002.sig")
            , testCase
                "000056-012.ring_trust"
                (testSerialization "000056-012.ring_trust")
            , testCase
                "000057-013.user_id"
                (testSerialization "000057-013.user_id")
            , testCase "000058-002.sig" (testSerialization "000058-002.sig")
            , testCase
                "000059-012.ring_trust"
                (testSerialization "000059-012.ring_trust")
            , testCase
                "000060-007.secret_subkey"
                (testSerialization "000060-007.secret_subkey")
            , testCase "000061-002.sig" (testSerialization "000061-002.sig")
            , testCase
                "000062-012.ring_trust"
                (testSerialization "000062-012.ring_trust")
            , testCase
                "000063-005.secret_key"
                (testSerialization "000063-005.secret_key")
            , testCase "000064-002.sig" (testSerialization "000064-002.sig")
            , testCase
                "000065-012.ring_trust"
                (testSerialization "000065-012.ring_trust")
            , testCase
                "000066-013.user_id"
                (testSerialization "000066-013.user_id")
            , testCase "000067-002.sig" (testSerialization "000067-002.sig")
            , testCase
                "000068-012.ring_trust"
                (testSerialization "000068-012.ring_trust")
            , testCase
                "000069-005.secret_key"
                (testSerialization "000069-005.secret_key")
            , testCase
                "000070-013.user_id"
                (testSerialization "000070-013.user_id")
            , testCase "000071-002.sig" (testSerialization "000071-002.sig")
            , testCase
                "000072-012.ring_trust"
                (testSerialization "000072-012.ring_trust")
            , testCase
                "000073-017.attribute"
                (testSerialization "000073-017.attribute")
            , testCase "000074-002.sig" (testSerialization "000074-002.sig")
            , testCase
                "000075-012.ring_trust"
                (testSerialization "000075-012.ring_trust")
            , testCase
                "000076-007.secret_subkey"
                (testSerialization "000076-007.secret_subkey")
            , testCase "000077-002.sig" (testSerialization "000077-002.sig")
            , testCase
                "000078-012.ring_trust"
                (testSerialization "000078-012.ring_trust")
            , testCase "pubring.gpg" (testSerialization "pubring.gpg")
            , testCase "secring.gpg" (testSerialization "secring.gpg")
            , testCase
                "compressedsig.gpg"
                (testSerialization "compressedsig.gpg")
            , testCase
                "compressedsig-zlib.gpg"
                (testSerialization "compressedsig-zlib.gpg")
            , testCase
                "compressedsig-bzip2.gpg"
                (testSerialization "compressedsig-bzip2.gpg")
            , testCase "onepass_sig" (testSerialization "onepass_sig")
            , testCase
                "uncompressed-ops-dsa.gpg"
                (testSerialization "uncompressed-ops-dsa.gpg")
            , testCase
                "uncompressed-ops-rsa.gpg"
                (testSerialization "uncompressed-ops-rsa.gpg")
            , testCase "simple.seckey" (testSerialization "simple.seckey")
            , testCase
                "v3-genericcert.sig"
                (testSerialization "v3-genericcert.sig")
            , testCase
                "sigs-with-regexes"
                (testSerialization "sigs-with-regexes")
            , testCase
                "gnu-dummy-s2k-101-secret-key.gpg"
                (testSerialization "gnu-dummy-s2k-101-secret-key.gpg")
            , testCase
                "anibal-ed25519.gpg"
                (testSerialization "anibal-ed25519.gpg")
            , testCase
                "nist_p-256_key.gpg"
                (testSerialization "nist_p-256_key.gpg")
            , testCase
                "nist_p-256_secretkey.gpg"
                (testSerialization "nist_p-256_secretkey.gpg")
            , testCase
                "v6-secret.pgp.aa"
                (testSerialization "v6-secret.pgp.aa")
            , testCase
                "sample-eddsa.pubkey"
                (testSerialization "sample-eddsa.pubkey")
            , testCase
                "should not serialize point at infinity"
                testPointAtInfinitySerialization
            , testCase
                "should reject mismatched EC coordinate widths"
                testPointSerializationRejectsMismatchedCoordinateWidths
            ]
        , testGroup
            "TKUnknown Serialization group"
            [ testCase
                "pubring.gpg TKUnknown serialization"
                (testTKSerialization "pubring.gpg")
            , testCase
                "secring.gpg TKUnknown serialization"
                (testTKSerialization "secring.gpg")
            ]
        , testGroup
            "Argon2 S2K group"
            [ testCase
                "Argon2 SKESK packet roundtrip"
                testArgon2S2KPacketRoundTrip
            ]
        , testGroup
            "RFC9580 SEIPD v2 group"
            [ testCase "SEIPD v2 packet roundtrip" testSEIPDv2PacketRoundTrip
            , testCase
                "SEIPD v2 rejects invalid chunk size"
                testSEIPDv2RejectInvalidChunkSize
            , testCase
                "PKESKv6 parsing does not fall back to legacy parser"
                testPKESKv6ParsesAsV6WithoutLegacyFallback
            , testCase
                "PKESKv6 rejects invalid recipient key version"
                testPKESKv6RejectsInvalidRecipientIdentifierVersion
            , testCase
                "PKESKv6 rejects recipient length/version mismatches"
                testPKESKv6RejectsRecipientIdentifierLengthVersionMismatch
            , testCase
                "legacy RSA PKESK rejects extra MPIs"
                testLegacyPKESKRSARejectsExtraMPI
            , testCase
                "legacy ECDH PKESK rejects wrong MPI count"
                testLegacyPKESKECDHRejectsWrongMPICount
            , testCase
                "legacy X25519 PKESK rejects wrong MPI count"
                testLegacyPKESKX25519RejectsWrongMPICount
            , testCase
                "legacy unencrypted secret key rejects checksum mismatch"
                testLegacyUnencryptedSecretKeyRejectsChecksumMismatch
            , testCase
                "legacy secret key rejects unsupported symmetric algorithm IV sizing"
                testLegacySecretKeyRejectsUnsupportedSymmetricAlgorithmIVSizing
            , testCase
                "legacy X25519 PKESK parses RFC9580 v3 octet layout"
                testLegacyPKESKX25519ParsesRFC9580V3OctetLayout
            , testCase
                "SigV6 rejects invalid salt size"
                testSigV6RejectsInvalidSaltSize
            , testCase
                "OPS6 rejects invalid salt size"
                testOPS6RejectsInvalidSaltSize
            , testCase
                "OPS3 rejects invalid nested-flag octet"
                testOPS3RejectsInvalidNestedFlagOctet
            , testCase
                "OPS6 rejects invalid nested-flag octet"
                testOPS6RejectsInvalidNestedFlagOctet
            , testCase
                "ECDH pubkey rejects reserved KDF length 0"
                testECDHPubkeyRejectsReservedKDFLengthZero
            , testCase
                "ECDH pubkey rejects reserved KDF length 255"
                testECDHPubkeyRejectsReservedKDFLength255
            , testCase
                "ECDH pubkey encodes fixed KDF length trailer"
                testECDHPubkeyEncodesFixedKDFLength
            , testCase
                "empty key-flags subpacket encodes explicit zero octet"
                testKeyFlagsSubpacketEncodesEmptySetWithZeroOctet
            , testCase
                "v6-secret fixture SigV6 semantics"
                testV6SecretFixtureSignatureSemantics
            , testCase
                "v6-secret fixture derives eight-octet key-id from fingerprint prefix"
                testV6SecretFixtureDerivesEightOctetKeyID
            ]
        ]

testSerialization :: FilePath -> Assertion
testSerialization fpr = do
    bs <- readFixturePayload fpr
    let firstpass = runGet get bs
    case fmap unBlock firstpass of
        Left _ -> assertFailure $ "First pass failed on " ++ fpr
        Right [] ->
            assertFailure $ "First pass of " ++ fpr ++ " decoded to nothing."
        Right packs -> do
            let roundtrip = runPut $ put (Block packs)
            let secondpass = runGet (get :: Get (Block Pkt)) roundtrip
            if fmap unBlock secondpass == Right []
                then
                    assertFailure $
                        "Second pass of " ++ fpr ++ " decoded to nothing."
                else assertEqual ("for " ++ fpr) firstpass secondpass

testTKSerialization :: FilePath -> Assertion
testTKSerialization fpr = do
    bs <- readFixturePayload fpr
    let pkts = parsePktsWithWireRep (wireRepRef bs) bs
        tksWithWireRep = parseTKsWithWireRep True pkts
    if null tksWithWireRep
        then
            assertFailure $
                "TKUnknown serialization test: " ++ fpr ++ " parsed to no TKs"
        else forM_ tksWithWireRep (testTKRoundtrip fpr)

testTKRoundtrip :: FilePath -> TKWithWireRep -> Assertion
testTKRoundtrip fpr tk = do
    let packets = _tkPackets tk
        encoded =
            runPut
                (put (Block (map (\p -> p ^. pktWireRep . pktValue) packets)))
    case runGet (get :: Get (Block Pkt)) encoded of
        Left err ->
            assertFailure $
                "TKUnknown " ++ fpr ++ " packet re-parse failed: " ++ err
        Right reparsedBlock ->
            assertEqual
                ("TKUnknown packet re-serialization roundtrip for " ++ fpr)
                (Block (map (\p -> p ^. pktWireRep . pktValue) packets))
                reparsedBlock
    case toStructuredTKWithWireRep tk of
        Left err ->
            assertFailure $
                "TKUnknown structured conversion failed for "
                    ++ fpr
                    ++ ": "
                    ++ show err
        Right structured ->
            case canonicalizeTKStructuredWithWireRep structured of
                Left err ->
                    assertFailure $
                        "TKUnknown canonical conversion failed for "
                            ++ fpr
                            ++ ": "
                            ++ show err
                Right _canonical -> pure ()

testArgon2S2KPacketRoundTrip :: Assertion
testArgon2S2KPacketRoundTrip = do
    let s2k = Argon2 (Salt16 (B.pack [0x00 .. 0x0f])) 1 4 15
        pkt =
            SKESKPkt
                (SKESKPayloadV4Packet (SKESKPayloadV4 AES128 s2k Nothing))
        encoded = runPut (put pkt)
    assertEqual
        "Argon2 S2K SKESK packet roundtrip"
        (Right pkt)
        (runGet (get :: Get Pkt) encoded)

testIssuerFingerprintRejectsUnknownVersion :: Assertion
testIssuerFingerprintRejectsUnknownVersion = do
    let encoded =
            runPut $ do
                putWord8 34
                putWord8 33
                putWord8 5
                putByteString (B.replicate 32 0)
    case runGet (get :: Get SigSubPacket) encoded of
        Left _ -> pure ()
        Right _ ->
            assertFailure
                "issuer fingerprint subpacket version 5 should be rejected"

testV6Ed25519PublicKeySerializesFixedLength :: Assertion
testV6Ed25519PublicKeySerializesFixedLength = do
    let pkp =
            PKPayload
                V6
                (ThirtyTwoBitTimeStamp 0)
                0
                EdDSA
                (EdDSAPubKey EdSigningCurve25519 (NativeEPoint (EPoint 1)))
        encoded = runPut (put (PublicKeyPkt pkp))
    assertEqual
        "v6 Ed25519 public key serialization uses fixed-length octet strings"
        44
        (BL.length encoded)

testV4Ed25519PublicKeyParsesNativeFixedLength :: Assertion
testV4Ed25519PublicKeyParsesNativeFixedLength = do
    let raw = B.pack [0x01 .. 0x20]
        encoded =
            runPut $ do
                putWord8 0xc6
                putWord8 38
                putWord8 4
                putWord32be 0
                putWord8 (fromIntegral (fromFVal PKA.Ed25519))
                putByteString raw
    case runGet (get :: Get Pkt) encoded of
        Right
            ( PublicKeyPkt
                    ( PKPayload
                            V4
                            (ThirtyTwoBitTimeStamp 0)
                            _
                            PKA.Ed25519
                            (EdDSAPubKey EdSigningCurve25519 (NativeEPoint (EPoint x)))
                        )
                ) ->
                assertEqual
                    "v4 Ed25519 fixed-length key material parsed as native point"
                    (os2ip raw)
                    x
        other ->
            assertFailure
                ( "Expected v4 Ed25519 fixed-length public key parse, got "
                    ++ show other
                )

testV4X25519PublicSubkeyParsesNativeFixedLength :: Assertion
testV4X25519PublicSubkeyParsesNativeFixedLength = do
    let raw = B.pack [0x01 .. 0x20]
        encoded =
            runPut $ do
                putWord8 0xce
                putWord8 38
                putWord8 4
                putWord32be 0
                putWord8 (fromIntegral (fromFVal PKA.X25519))
                putByteString raw
    case runGet (get :: Get Pkt) encoded of
        Right
            ( PublicSubkeyPkt
                    ( PKPayload
                            V4
                            (ThirtyTwoBitTimeStamp 0)
                            _
                            PKA.X25519
                            (EdDSAPubKey EdSigningCurve25519 (NativeEPoint (EPoint x)))
                        )
                ) ->
                assertEqual
                    "v4 X25519 fixed-length key material parsed as native point"
                    (os2ip raw)
                    x
        other ->
            assertFailure
                ( "Expected v4 X25519 fixed-length public subkey parse, got "
                    ++ show other
                )

testV6SignatureIssuerFingerprintVersionMismatch :: Assertion
testV6SignatureIssuerFingerprintVersionMismatch = do
    let signer =
            PKPayload
                V6
                (ThirtyTwoBitTimeStamp 0)
                0
                EdDSA
                (EdDSAPubKey EdSigningCurve25519 (NativeEPoint (EPoint 1)))
        sig =
            SignaturePkt
                ( SigV6
                    BinarySig
                    EdDSA
                    SHA512
                    (SignatureSalt (BL.replicate 32 0))
                    [ SigSubPacket
                        False
                        (IssuerFingerprint IssuerFingerprintV4 (fingerprint signer))
                    ]
                    []
                    0
                    (NE.fromList [MPI 0, MPI 0])
                )
        verifier _ _ _ =
            Right
                ( Verification
                    signer
                    (case sig of SignaturePkt sp -> sp; _ -> error "impossible")
                    []
                )
    case verifySigWith
        defaultVerificationPolicy
        verifier
        sig
        emptyPSC
        Nothing of
        Left IssuerFingerprintSubpacketMismatch -> pure ()
        Left err -> assertFailure ("unexpected verification error: " ++ show err)
        Right _ ->
            assertFailure
                "v6 signature with issuer fingerprint version 4 should be rejected"

testPointAtInfinitySerialization :: Assertion
testPointAtInfinitySerialization =
    assertEqual
        "point at infinity should not serialize"
        Nothing
        (point2MBS ECCT.PointO)

testPointSerializationRejectsMismatchedCoordinateWidths
    :: Assertion
testPointSerializationRejectsMismatchedCoordinateWidths =
    assertEqual
        "point serialization should reject mismatched coordinate widths"
        Nothing
        (point2MBS (ECCT.Point 1 256))

testSEIPDv2PacketRoundTrip :: Assertion
testSEIPDv2PacketRoundTrip = do
    let salt = Salt (B.pack [0x00 .. 0x1f])
        pkt =
            SymEncIntegrityProtectedDataPkt
                (SEIPD2 AES256 OCB 16 salt "\x01\x02\x03\x04")
        encoded = runPut (put pkt)
    assertEqual
        "SEIPD v2 packet roundtrip"
        (Right pkt)
        (runGet (get :: Get Pkt) encoded)

testSEIPDv2RejectInvalidChunkSize :: Assertion
testSEIPDv2RejectInvalidChunkSize = do
    let encoded =
            runPut $ do
                putWord8 0xd2
                putWord8 37
                putWord8 2
                putWord8 7
                putWord8 2
                putWord8 17
                putByteString (B.replicate 32 0)
                putWord8 0
    case runGet (get :: Get Pkt) encoded of
        Right BrokenPacketPkt {} -> return ()
        other ->
            assertFailure
                ( "SEIPD v2 parser should reject chunk sizes larger than 16, got "
                    ++ show other
                )

testPKESKv6ParsesAsV6WithoutLegacyFallback :: Assertion
testPKESKv6ParsesAsV6WithoutLegacyFallback = do
    let recipientKeyIdentifier = BL.pack (0x04 : replicate 20 0)
        esk = "\x00\x00"
        encoded =
            runPut $ do
                putWord8 0xc1
                putWord8 26
                putWord8 6
                putWord8 21
                putByteString (BL.toStrict recipientKeyIdentifier)
                putWord8 1
                putByteString (BL.toStrict esk)
    case runGet (get :: Get Pkt) encoded of
        Right
            ( PKESKPkt
                    (PKESKPayloadV6Packet (PKESKPayloadV6 rid pka parsedEsk))
                ) -> do
                assertEqual
                    "PKESKv6 recipient key identifier"
                    recipientKeyIdentifier
                    rid
                assertEqual "PKESKv6 algorithm" RSA pka
                assertEqual "PKESKv6 ESK payload" esk parsedEsk
        other ->
            assertFailure
                ("Expected PKESKPkt (PKESK6 ...) parse result, got " ++ show other)

testPKESKv6RejectsInvalidRecipientIdentifierVersion :: Assertion
testPKESKv6RejectsInvalidRecipientIdentifierVersion = do
    let recipientKeyIdentifier = BL.pack (0x05 : replicate 20 0)
        encoded =
            runPut $ do
                putWord8 0xc1
                putWord8 24
                putWord8 6
                putWord8 21
                putByteString (BL.toStrict recipientKeyIdentifier)
                putWord8 1
    case runGet (get :: Get Pkt) encoded of
        Right BrokenPacketPkt {} -> pure ()
        other ->
            assertFailure
                ( "Expected malformed PKESKv6 key version to produce BrokenPacketPkt, got "
                    ++ show other
                )

testPKESKv6RejectsRecipientIdentifierLengthVersionMismatch
    :: Assertion
testPKESKv6RejectsRecipientIdentifierLengthVersionMismatch = do
    let recipientKeyIdentifier = BL.pack (0x04 : replicate 32 0)
        encoded =
            runPut $ do
                putWord8 0xc1
                putWord8 36
                putWord8 6
                putWord8 33
                putByteString (BL.toStrict recipientKeyIdentifier)
                putWord8 1
    case runGet (get :: Get Pkt) encoded of
        Right BrokenPacketPkt {} -> pure ()
        other ->
            assertFailure
                ( "Expected malformed PKESKv6 recipient length/version mismatch to produce BrokenPacketPkt, got "
                    ++ show other
                )

testLegacyPKESKRSARejectsExtraMPI :: Assertion
testLegacyPKESKRSARejectsExtraMPI = do
    let encoded =
            runPut $ do
                putWord8 0xc1
                putWord8 16
                putWord8 3
                putByteString (B.replicate 8 0)
                putWord8 1
                put (MPI 1)
                put (MPI 2)
    case runGet (get :: Get Pkt) encoded of
        Right BrokenPacketPkt {} -> pure ()
        other ->
            assertFailure
                ( "Expected malformed legacy RSA PKESK to produce BrokenPacketPkt, got "
                    ++ show other
                )

testLegacyPKESKECDHRejectsWrongMPICount :: Assertion
testLegacyPKESKECDHRejectsWrongMPICount = do
    let encoded =
            runPut $ do
                putWord8 0xc1
                putWord8 13
                putWord8 3
                putByteString (B.replicate 8 0)
                putWord8 18
                put (MPI 1)
    case runGet (get :: Get Pkt) encoded of
        Right BrokenPacketPkt {} -> pure ()
        other ->
            assertFailure
                ( "Expected malformed legacy ECDH PKESK to produce BrokenPacketPkt, got "
                    ++ show other
                )

testLegacyPKESKX25519RejectsWrongMPICount :: Assertion
testLegacyPKESKX25519RejectsWrongMPICount = do
    let encoded =
            runPut $ do
                putWord8 0xc1
                putWord8 13
                putWord8 3
                putByteString (B.replicate 8 0)
                putWord8 (fromIntegral (fromFVal X25519))
                put (MPI 1)
    case runGet (get :: Get Pkt) encoded of
        Right BrokenPacketPkt {} -> pure ()
        other ->
            assertFailure
                ( "Expected malformed legacy X25519 PKESK to produce BrokenPacketPkt, got "
                    ++ show other
                )

testLegacyUnencryptedSecretKeyRejectsChecksumMismatch
    :: Assertion
testLegacyUnencryptedSecretKeyRejectsChecksumMismatch = do
    secretPackets <-
        DC.runConduitRes $
            CB.sourceFile "tests/data/unencrypted.seckey"
                DC..| conduitGet get
                DC..| CL.consume
    pkt <-
        case secretPackets of
            (SecretKeyPkt pkp (SUUnencrypted sk checksum) : _) ->
                pure (SecretKeyPkt pkp (SUUnencrypted sk (checksum `xor` 1)))
            (SecretKeyPkt _ _ : _) ->
                assertFailure
                    "unencrypted.seckey did not begin with an unencrypted secret key packet"
                    >> fail "expected unencrypted secret key packet"
            _ ->
                assertFailure
                    "unencrypted.seckey did not begin with a secret key packet"
                    >> fail "expected secret key packet"
    let encoded = runPut (put pkt)
    case runGet (get :: Get Pkt) encoded of
        Right BrokenPacketPkt {} -> pure ()
        other ->
            assertFailure
                ( "Expected unencrypted secret key checksum mismatch to produce BrokenPacketPkt, got "
                    ++ show other
                )

testLegacySecretKeyRejectsUnsupportedSymmetricAlgorithmIVSizing
    :: Assertion
testLegacySecretKeyRejectsUnsupportedSymmetricAlgorithmIVSizing = do
    let pkt =
            SecretKeyPkt
                (PKPayload V4 0 0 RSA (UnknownPKey BL.empty))
                ( SUSSHA1
                    (OtherSA 0xfe)
                    (IteratedSalted SHA256 (Salt8 "12345678") (IterationCount 65536))
                    (IV (B.replicate 8 0))
                    BL.empty
                )
        encoded = runPut (put pkt)
    case runGet (get :: Get Pkt) encoded of
        Right BrokenPacketPkt {} -> pure ()
        other ->
            assertFailure
                ( "Expected legacy secret key with unsupported symmetric algorithm to produce BrokenPacketPkt, got "
                    ++ show other
                )

testLegacyPKESKX25519ParsesRFC9580V3OctetLayout :: Assertion
testLegacyPKESKX25519ParsesRFC9580V3OctetLayout = do
    let ephemeral = B.pack [0x01 .. 0x20]
        wrappedWithAlgo =
            B.singleton (fromIntegral (fromFVal AES128))
                <> B.replicate 24 0x5a
        encoded =
            runPut $ do
                putWord8 0xc1
                putWord8
                    ( fromIntegral
                        (1 + 8 + 1 + B.length ephemeral + 1 + B.length wrappedWithAlgo)
                    )
                putWord8 3
                putByteString (B.replicate 8 0)
                putWord8 (fromIntegral (fromFVal X25519))
                putByteString ephemeral
                putWord8 (fromIntegral (B.length wrappedWithAlgo))
                putByteString wrappedWithAlgo
    case runGet (get :: Get Pkt) encoded of
        Right
            ( PKESKPkt
                    ( PKESKPayloadV3Packet
                            (PKESKPayloadV3 3 _ X25519 (ephMPI :| [eskMPI]))
                        )
                ) -> do
                assertEqual
                    "legacy X25519 v3 octet-layout ephemeral parse"
                    (MPI (os2ip ephemeral))
                    ephMPI
                assertEqual
                    "legacy X25519 v3 octet-layout wrapped parse"
                    (MPI (os2ip wrappedWithAlgo))
                    eskMPI
        other ->
            assertFailure
                ( "Expected legacy X25519 PKESK v3 octet-layout parse success, got "
                    ++ show other
                )

testSigV6RejectsInvalidSaltSize :: Assertion
testSigV6RejectsInvalidSaltSize = do
    let encoded =
            runPut $ do
                putWord8 0xc2
                putWord8 45
                putWord8 6
                putWord8 0
                putWord8 1
                putWord8 8
                putWord32be 0
                putWord32be 0
                putWord16be 0
                putWord8 32
                putByteString (B.replicate 32 0)
                putWord16be 0
    case runGet (get :: Get Pkt) encoded of
        Right BrokenPacketPkt {} -> return ()
        other ->
            assertFailure
                ( "Expected invalid SigV6 salt size to produce BrokenPacketPkt, got "
                    ++ show other
                )

testOPS6RejectsInvalidSaltSize :: Assertion
testOPS6RejectsInvalidSaltSize = do
    let encoded =
            runPut $ do
                putWord8 0xc4
                putWord8 69
                putWord8 6
                putWord8 0
                putWord8 10
                putWord8 22
                putWord8 31
                putByteString (B.replicate 31 0)
                putByteString (B.replicate 32 0)
                putWord8 0
    case runGet (get :: Get Pkt) encoded of
        Right BrokenPacketPkt {} -> return ()
        other ->
            assertFailure
                ( "Expected invalid OPS6 salt size to produce BrokenPacketPkt, got "
                    ++ show other
                )

testOPS3RejectsInvalidNestedFlagOctet :: Assertion
testOPS3RejectsInvalidNestedFlagOctet = do
    let encoded =
            runPut $ do
                putWord8 0xc4
                putWord8 13
                putWord8 3
                putWord8 0
                putWord8 8
                putWord8 1
                putByteString (B.replicate 8 0)
                putWord8 2
    case runGet (get :: Get Pkt) encoded of
        Right BrokenPacketPkt {} -> pure ()
        other ->
            assertFailure
                ( "Expected invalid OPS3 nested-flag octet to produce BrokenPacketPkt, got "
                    ++ show other
                )

testOPS6RejectsInvalidNestedFlagOctet :: Assertion
testOPS6RejectsInvalidNestedFlagOctet = do
    let encoded =
            runPut $ do
                putWord8 0xc4
                putWord8 70
                putWord8 6
                putWord8 0
                putWord8 8
                putWord8 1
                putWord8 32
                putByteString (B.replicate 32 0)
                putByteString (B.replicate 32 0)
                putWord8 2
    case runGet (get :: Get Pkt) encoded of
        Right BrokenPacketPkt {} -> pure ()
        other ->
            assertFailure
                ( "Expected invalid OPS6 nested-flag octet to produce BrokenPacketPkt, got "
                    ++ show other
                )

mkECDHBoundaryTestPacket :: Pkt
mkECDHBoundaryTestPacket =
    PublicKeyPkt
        ( PKPayload
            V4
            0
            0
            ECDH
            ( ECDHPubKey
                ( ECDSAPubKey
                    ( ECDSA_PublicKey
                        ( ECDSA.PublicKey
                            (ECCT.getCurveByName ECCT.SEC_p256r1)
                            (ECCT.Point 1 2)
                        )
                    )
                )
                SHA256
                AES128
            )
        )

setStrictByteAt
    :: Int -> Word8 -> B.ByteString -> Maybe B.ByteString
setStrictByteAt idx w bs
    | idx < 0 || idx >= B.length bs = Nothing
    | otherwise =
        let (prefix, rest) = B.splitAt idx bs
         in case B.uncons rest of
                Nothing -> Nothing
                Just (_, suffix) -> Just (prefix <> B.singleton w <> suffix)

testECDHPubkeyRejectsReservedKDFLengthZero :: Assertion
testECDHPubkeyRejectsReservedKDFLengthZero = do
    let encoded = BL.toStrict (runPut (put mkECDHBoundaryTestPacket))
    mutated <-
        maybe
            ( assertFailure
                "failed to locate ECDH KDF length byte for zero-length rejection test"
                >> fail "unreachable"
            )
            pure
            (setStrictByteAt (B.length encoded - 4) 0x00 encoded)
    case runGet (get :: Get Pkt) (BL.fromStrict mutated) of
        Right BrokenPacketPkt {} -> pure ()
        other ->
            assertFailure
                ( "Expected ECDH KDF length 0 to produce BrokenPacketPkt, got "
                    ++ show other
                )

testECDHPubkeyRejectsReservedKDFLength255 :: Assertion
testECDHPubkeyRejectsReservedKDFLength255 = do
    let encoded = BL.toStrict (runPut (put mkECDHBoundaryTestPacket))
    mutated <-
        maybe
            ( assertFailure
                "failed to locate ECDH KDF length byte for 0xff rejection test"
                >> fail "unreachable"
            )
            pure
            (setStrictByteAt (B.length encoded - 4) 0xff encoded)
    case runGet (get :: Get Pkt) (BL.fromStrict mutated) of
        Right BrokenPacketPkt {} -> pure ()
        other ->
            assertFailure
                ( "Expected ECDH KDF length 255 to produce BrokenPacketPkt, got "
                    ++ show other
                )

testECDHPubkeyEncodesFixedKDFLength :: Assertion
testECDHPubkeyEncodesFixedKDFLength = do
    let encoded = BL.toStrict (runPut (put mkECDHBoundaryTestPacket))
        trailer = B.drop (B.length encoded - 4) encoded
    assertEqual
        "ECDH public-key encoding should emit fixed KDF trailer [3,1,hash,sym]"
        (B.pack [0x03, 0x01, fromFVal SHA256, fromFVal AES128])
        trailer

testKeyFlagsSubpacketEncodesEmptySetWithZeroOctet :: Assertion
testKeyFlagsSubpacketEncodesEmptySetWithZeroOctet = do
    let encoded = runPut (put (SigSubPacket False (KeyFlags Set.empty)))
    assertEqual
        "empty key-flags subpacket should encode an explicit zero flags octet"
        [2, 27, 0]
        (BL.unpack encoded)
    case runGet (get :: Get SigSubPacket) encoded of
        Right (SigSubPacket False (KeyFlags flags)) ->
            assertEqual
                "empty key-flags subpacket should decode back to an empty flag set"
                Set.empty
                flags
        other ->
            assertFailure
                ( "Expected empty key-flags subpacket roundtrip, got "
                    ++ show other
                )

expectedV6SaltSizeForTest :: HashAlgorithm -> Maybe Int
expectedV6SaltSizeForTest =
    fmap fromIntegral . signatureV6SaltSizeForHashAlgorithm

signatureHasIssuerFingerprintV6
    :: Fingerprint -> SignaturePayload -> Bool
signatureHasIssuerFingerprintV6 expectedFp (SigV6 _ _ _ _ hashed unhashed _ _) =
    expectedFp
        `elem` [ ifp
               | SigSubPacket _ (IssuerFingerprint IssuerFingerprintV6 ifp) <-
                    hashed ++ unhashed
               ]
signatureHasIssuerFingerprintV6 _ _ = False

testV6SecretFixtureSignatureSemantics :: Assertion
testV6SecretFixtureSignatureSemantics = do
    armors <- loadArmor "v6-secret.pgp.aa"
    payload <-
        case armors of
            (a : _) -> pure (armorPayload a)
            [] ->
                assertFailure
                    "v6-secret.pgp.aa should contain one armored payload"
                    >> fail "expected one armored payload"
    let packets = parsePkts payload
        primaryV6Key =
            listToMaybe
                [ pkp
                | SecretKeyPkt pkp _ <- packets
                , _keyVersion pkp == V6
                ]
                <|> listToMaybe
                    [ pkp
                    | PublicKeyPkt pkp <- packets
                    , _keyVersion pkp == V6
                    ]
        signatures =
            [ (sig, ha, salt)
            | SignaturePkt sig@(SigV6 _ _ ha salt _ _ _ _) <- packets
            ]
    pkp <-
        case primaryV6Key of
            Nothing ->
                assertFailure
                    "v6-secret.pgp.aa should contain a primary v6 key packet"
                    >> fail "expected primary v6 key packet"
            Just k -> pure k
    assertBool
        "v6-secret.pgp.aa should contain at least one SigV6 packet"
        (not (null signatures))
    mapM_
        ( \(_, ha, salt) ->
            case expectedV6SaltSizeForTest ha of
                Nothing ->
                    assertFailure
                        ( "SigV6 in v6-secret.pgp.aa uses unsupported salt hash algorithm: "
                            ++ show ha
                        )
                Just expected ->
                    assertEqual
                        "SigV6 salt size in v6-secret.pgp.aa should match hash algorithm"
                        expected
                        (fromIntegral (BL.length (unSignatureSalt salt)))
        )
        signatures
    assertBool
        "v6-secret.pgp.aa should include at least one IssuerFingerprint v6 matching the primary key"
        ( any
            ( \(sig, _, _) -> signatureHasIssuerFingerprintV6 (fingerprint pkp) sig
            )
            signatures
        )

testV6SecretFixtureDerivesEightOctetKeyID :: Assertion
testV6SecretFixtureDerivesEightOctetKeyID = do
    armors <- loadArmor "v6-secret.pgp.aa"
    payload <-
        case armors of
            (a : _) -> pure (armorPayload a)
            [] ->
                assertFailure
                    "v6-secret.pgp.aa should contain one armored payload"
                    >> fail "expected one armored payload"
    let packets = parsePkts payload
        primaryV6Key =
            listToMaybe
                [ pkp
                | SecretKeyPkt pkp _ <- packets
                , _keyVersion pkp == V6
                ]
                <|> listToMaybe
                    [ pkp
                    | PublicKeyPkt pkp <- packets
                    , _keyVersion pkp == V6
                    ]
    pkp <-
        case primaryV6Key of
            Nothing ->
                assertFailure
                    "v6-secret.pgp.aa should contain a primary v6 key packet"
                    >> fail "expected primary v6 key packet"
            Just k -> pure k
    derivedKeyId <-
        case eightOctetKeyID pkp of
            Left err ->
                assertFailure
                    ("Expected v6 eight-octet key-id derivation to succeed: " ++ err)
                    >> fail "expected v6 eight-octet key-id"
            Right keyId -> pure keyId
    let expectedKeyId = EightOctetKeyId (BL.take 8 (unFingerprint (fingerprint pkp)))
    assertEqual
        "v6 eight-octet key-id should be the high-order 64 bits of the fingerprint"
        expectedKeyId
        derivedKeyId
