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 Codec.Encryption.OpenPGP.Types
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 ()

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
         ] ++
         [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))
      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
        ] &&
        (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 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