Arbitrary.hs

-- Arbitrary.hs: QuickCheck instances
-- Copyright © 2014-2026  Clint Adams
-- This software is released under the terms of the Expat license.
-- (See the LICENSE file).
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleInstances #-}

module Codec.Encryption.OpenPGP.Arbitrary
    (
    ) where

import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL
import qualified Data.List.NonEmpty as NE
import Data.Maybe (fromMaybe)
import Network.URI (nullURI, parseURI)
import Test.QuickCheck
    ( Arbitrary (..)
    , Gen
    , choose
    , elements
    , frequency
    , getPositive
    , listOf1
    , oneof
    , sized
    , suchThat
    , vector
    )
import Test.QuickCheck.Instances ()

import Codec.Encryption.OpenPGP.Types

instance Arbitrary (PKESK 'PKESKV3) where
    arbitrary = do
        eoki <- arbitrary
        pka <- arbitrary
        mpis <-
            case pka of
                RSA -> (\mpi -> mpi NE.:| []) <$> arbitrary
                DeprecatedRSAEncryptOnly -> (\mpi -> mpi NE.:| []) <$> arbitrary
                ElgamalEncryptOnly -> do
                    mpi1 <- arbitrary
                    mpi2 <- arbitrary
                    pure (mpi1 NE.:| [mpi2])
                ForbiddenElgamal -> do
                    mpi1 <- arbitrary
                    mpi2 <- arbitrary
                    pure (mpi1 NE.:| [mpi2])
                ECDH -> do
                    mpi1 <- arbitrary
                    mpi2 <- arbitrary
                    pure (mpi1 NE.:| [mpi2])
                _ -> arbitrary
        pure (PKESK3Packet 3 eoki pka mpis)

instance Arbitrary (PKESK 'PKESKV6) where
    arbitrary = do
        rid <-
            oneof
                [ pure BL.empty
                , BL.pack <$> vector 20
                , BL.pack . (4 :) <$> vector 20
                , BL.pack <$> vector 32
                , BL.pack . (6 :) <$> vector 32
                ]
        pka <- arbitrary
        esk <- BL.pack <$> listOf1 arbitrary
        pure (PKESK6Packet rid pka esk)

instance Arbitrary (SKESK 'SKESKV4) where
    arbitrary = do
        sa <- elements [AES128, AES192, AES256]
        s2k <- arbitrarySKESKv4S2K
        esk <- oneof [pure Nothing, Just . BL.pack <$> listOf1 arbitrary]
        pure (SKESK4Packet sa s2k esk)

instance Arbitrary (SKESK 'SKESKV6) where
    arbitrary = do
        sa <- elements [AES128, AES192, AES256]
        aead <- elements [EAX, OCB, GCM]
        s2k <- arbitrarySKESKv6S2K
        iv <- BL.pack <$> vector 16
        esk <- BL.pack <$> listOf1 arbitrary
        tag <- BL.pack <$> vector 16
        pure (SKESK6Packet sa aead s2k iv esk tag)

arbitrarySKESKv4S2K :: Gen S2K
arbitrarySKESKv4S2K =
    oneof
        [ Simple <$> elements supportedSKESKHashAlgorithms
        , Salted
            <$> elements supportedSKESKHashAlgorithms
            <*> (Salt8 . B.pack <$> vector 8)
        ]

arbitrarySKESKv6S2K :: Gen S2K
arbitrarySKESKv6S2K =
    Argon2
        <$> (Salt16 . B.pack <$> vector 16)
        <*> choose (1, 4)
        <*> choose (1, 4)
        <*> choose (3, 31)

supportedSKESKHashAlgorithms :: [HashAlgorithm]
supportedSKESKHashAlgorithms = [SHA1, SHA256, SHA384, SHA512]

instance Arbitrary Signature where
    arbitrary = fmap Signature arbitrary

instance Arbitrary UserId where
    arbitrary = fmap UserId arbitrary

--
instance Arbitrary SignaturePayload where
    arbitrary = frequency [(2, three), (3, four), (1, other)]
      where
        three = do
            st <- arbitrary
            w32 <- arbitrary
            eoki <- arbitrary
            pka <- arbitrary
            ha <- arbitrary
            w16 <- arbitrary
            SigV3 st w32 eoki pka ha w16 <$> arbitrary
        four = do
            st <- arbitrary
            pka <- arbitrary
            ha <- arbitrary
            has <- arbitrary
            uhas <- arbitrary
            w16 <- arbitrary
            SigV4 st pka ha has uhas w16 <$> arbitrary
        other = do
            v <- oneof [pure 5, choose (7, maxBound)]
            SigVOther v <$> arbitrary

instance Arbitrary SigSubPacket where
    arbitrary = do
        crit <- arbitrary
        SigSubPacket crit <$> arbitrary

instance Arbitrary SigSubPacketPayload where
    arbitrary =
        sized $ \n ->
            oneof
                ( [ sct
                  , set
                  , ec
                  , ts
                  , re
                  , ket
                  , psa
                  , rk
                  , i
                  , nd
                  , phas
                  , pcas
                  , ksps
                  , pks
                  , puid
                  , purl
                  , kfs
                  , suid
                  , rfr
                  , fs
                  , st
                  , udss
                  , oss
                  , ifp
                  , irfp
                  , pacs
                  ]
                    ++ [es n | n > 0]
                )
      where
        sct = fmap SigCreationTime arbitrary
        set = fmap SigExpirationTime arbitrary
        ec = fmap ExportableCertification arbitrary
        ts =
            arbitrary >>= \tl -> arbitrary >>= \ta -> return (TrustSignature tl ta)
        re = fmap RegularExpression arbitrary
        ket = fmap KeyExpirationTime arbitrary
        psa = fmap PreferredSymmetricAlgorithms arbitrary
        rk =
            arbitrary >>= \rcs ->
                arbitrary >>= \pka ->
                    fmap (RevocationKey rcs pka . Fingerprint . BL.pack) (vector 20)
        i = fmap Issuer arbitrary
        nd =
            arbitrary >>= \nfs ->
                arbitrary >>= \nn ->
                    arbitrary >>= \nv -> return (NotationData nfs nn nv)
        phas = fmap PreferredHashAlgorithms arbitrary
        pcas = fmap PreferredCompressionAlgorithms arbitrary
        ksps = fmap KeyServerPreferences arbitrary
        pks = fmap PreferredKeyServer arbitrary
        puid = fmap PrimaryUserId arbitrary
        purl =
            fmap (PolicyURL . URL . fromMaybe nullURI . parseURI) arbitrary
        kfs = fmap KeyFlags arbitrary
        suid = fmap SignersUserId arbitrary
        rfr =
            arbitrary >>= \rc ->
                arbitrary >>= \rr -> return (ReasonForRevocation rc rr)
        fs = fmap Features arbitrary
        st =
            arbitrary >>= \pka ->
                arbitrary >>= \ha ->
                    arbitrary >>= \sh -> return (SignatureTarget pka ha sh)
        es _ = pure (EmbeddedSignature (SigVOther 5 BL.empty))
        ifp =
            elements [IssuerFingerprintV4, IssuerFingerprintV6] >>= \v ->
                fmap
                    (IssuerFingerprint v)
                    ( if v == IssuerFingerprintV6
                        then fmap (Fingerprint . BL.pack) (vector 32)
                        else fmap (Fingerprint . BL.pack) (vector 20)
                    )
        irfp =
            elements [IssuerFingerprintV4, IssuerFingerprintV6] >>= \v ->
                fmap
                    (IntendedRecipient v)
                    ( if v == IssuerFingerprintV6
                        then fmap (Fingerprint . BL.pack) (vector 32)
                        else fmap (Fingerprint . BL.pack) (vector 20)
                    )
        pacs = fmap PreferredAEADCiphersuites arbitrary
        udss =
            choose (100, 110) >>= \a ->
                arbitrary >>= \b -> return (UserDefinedSigSub a b)
        oss =
            suchThat arbitrary isOtherSigSubType >>= \a ->
                arbitrary >>= \b -> return (OtherSigSub a b)
        isOtherSigSubType a =
            a <= 127
                && a
                    `notElem` [ 2
                              , 3
                              , 4
                              , 5
                              , 6
                              , 7
                              , 9
                              , 11
                              , 12
                              , 16
                              , 20
                              , 21
                              , 22
                              , 23
                              , 24
                              , 25
                              , 26
                              , 27
                              , 28
                              , 29
                              , 30
                              , 31
                              , 32
                              , 33
                              , 35
                              , 39
                              ]
                && (a < 100 || a > 110)

--
instance Arbitrary PubKeyAlgorithm where
    arbitrary = elements [RSA, DSA, ECDH, ECDSA, DH, EdDSA]

instance Arbitrary EightOctetKeyId where
    arbitrary = fmap (EightOctetKeyId . BL.pack) (vector 8)

instance Arbitrary Fingerprint where
    arbitrary =
        oneof
            [ fmap (Fingerprint . BL.pack) (vector 16) -- v3
            , fmap (Fingerprint . BL.pack) (vector 20) -- v4
            , fmap (Fingerprint . BL.pack) (vector 32) -- v6
            ]

instance Arbitrary MPI where
    arbitrary = fmap (MPI . getPositive) arbitrary

instance Arbitrary SigType where
    arbitrary =
        elements
            [ BinarySig
            , CanonicalTextSig
            , StandaloneSig
            , GenericCert
            , PersonaCert
            , CasualCert
            , PositiveCert
            , SubkeyBindingSig
            , PrimaryKeyBindingSig
            , SignatureDirectlyOnAKey
            , KeyRevocationSig
            , SubkeyRevocationSig
            , CertRevocationSig
            , TimestampSig
            , ThirdPartyConfirmationSig
            ]

instance Arbitrary HashAlgorithm where
    arbitrary =
        elements
            [ DeprecatedMD5
            , SHA1
            , RIPEMD160
            , SHA256
            , SHA384
            , SHA512
            , SHA224
            , SHA3_256
            , SHA3_512
            ]

instance Arbitrary SymmetricAlgorithm where
    arbitrary =
        elements
            [ Plaintext
            , IDEA
            , TripleDES
            , CAST5
            , Blowfish
            , ReservedSAFER
            , ReservedDES
            , AES128
            , AES192
            , AES256
            , Twofish
            , Camellia128
            , Camellia192
            , Camellia256
            ]

instance Arbitrary AEADAlgorithm where
    arbitrary = elements [EAX, OCB, GCM]

instance Arbitrary RevocationClass where
    arbitrary = frequency [(9, srk), (1, rco)]
      where
        srk = return SensitiveRK
        rco = fmap mkRevocationClass (choose (2, 6))

instance Arbitrary NotationFlag where
    arbitrary = frequency [(9, hr), (1, onf)]
      where
        hr = return HumanReadable
        onf = fmap mkNotationFlag (choose (1, 15))

instance Arbitrary CompressionAlgorithm where
    arbitrary = elements [Uncompressed, ZIP, ZLIB, BZip2]

instance Arbitrary KSPFlag where
    arbitrary = frequency [(9, nm), (1, kspo)]
      where
        nm = return NoModify
        kspo = fmap KSPOther (choose (2, 63))

instance Arbitrary KeyFlag where
    arbitrary =
        elements
            [ GroupKey
            , AuthKey
            , SplitKey
            , EncryptStorageKey
            , EncryptCommunicationsKey
            , SignDataKey
            , CertifyKeysKey
            ]

instance Arbitrary RevocationCode where
    arbitrary =
        elements
            [ NoReason
            , KeySuperseded
            , KeyMaterialCompromised
            , KeyRetiredAndNoLongerUsed
            , UserIdInfoNoLongerValid
            ]

instance Arbitrary FeatureFlag where
    arbitrary = frequency [(8, seipdV1), (1, seipdV2), (1, fo)]
      where
        seipdV1 = return FeatureSEIPDv1
        seipdV2 = return FeatureSEIPDv2
        fo =
            fmap FeatureOther (suchThat (choose (0, 63)) (`notElem` [4, 7]))

instance Arbitrary ThirtyTwoBitTimeStamp where
    arbitrary = fmap ThirtyTwoBitTimeStamp arbitrary

instance Arbitrary ThirtyTwoBitDuration where
    arbitrary = fmap ThirtyTwoBitDuration arbitrary

instance Arbitrary NotationName where
    arbitrary = fmap NotationName arbitrary

instance Arbitrary NotationValue where
    arbitrary = fmap NotationValue arbitrary

instance Arbitrary Padding where
    arbitrary = fmap Padding arbitrary