content: -- Copyright (C) 2015, 2016  Fraser Tweedale
  --
  -- Licensed under the Apache License, Version 2.0 (the "License");
  -- you may not use this file except in compliance with the License.
  -- You may obtain a copy of the License at
  --
  --      http://www.apache.org/licenses/LICENSE-2.0
  --
  -- Unless required by applicable law or agreed to in writing, software
  -- distributed under the License is distributed on an "AS IS" BASIS,
  -- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  -- See the License for the specific language governing permissions and
  -- limitations under the License.
  
  {-# LANGUAGE NoMonomorphismRestriction #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# OPTIONS_GHC -fno-warn-incomplete-uni-patterns #-}
  {-# OPTIONS_GHC -fno-warn-orphans #-}
  
  module Properties
    ( properties
    ) where
  
  import Control.Monad.IO.Class
  
  import Control.Lens ((&), set, view)
  import Crypto.Number.Basic (log2)
  import Crypto.Random
  import Data.Aeson (FromJSON, ToJSON, decode, encode)
  import qualified Data.ByteString as B
  
  import Hedgehog
  import qualified Hedgehog.Gen as Gen
  import qualified Hedgehog.Range as Range
  import Test.Tasty
  import Test.Tasty.Hedgehog
  
  import Crypto.JOSE.Types
  import Crypto.JOSE.JWK
  import Crypto.JOSE.JWS
  
  
  instance (MonadIO m) => MonadRandom (PropertyT m) where
    getRandomBytes = liftIO . getRandomBytes
  
  properties :: TestTree
  properties = testGroup "Properties"
    [ let n = "SizedBase64Integer round-trip" in testPropertyNamed n n (prop_roundTrip genSizedBase64Integer)
    , let n = "JWK round-trip" in testPropertyNamed n n (prop_roundTrip genJWK')
    , let n = "RSA gen, sign and verify" in testPropertyNamed n n prop_rsaSignAndVerify
    , let n = "gen, sign with best alg, verify" in testPropertyNamed n n prop_bestJWSAlg
    ]
  
  genBigInteger :: Gen Integer
  genBigInteger = Gen.integral $ Range.exponential 0 (2 ^ (4096 :: Integer))
  
  genBase64Integer :: Gen Base64Integer
  genBase64Integer = Base64Integer <$> genBigInteger
  
  genSizedBase64Integer :: Gen SizedBase64Integer
  genSizedBase64Integer = do
    x <- genBigInteger
    l <- Gen.element [0, 1, 2]  -- number of leading zero-bytes
    pure $ SizedBase64Integer ((log2 x `div` 8) + 1 + l) x
  
  
  prop_roundTrip :: (Eq a, Show a, ToJSON a, FromJSON a) => Gen a -> Property
  prop_roundTrip gen = property $
    forAll gen >>= \a -> decode (encode [a]) === Just [a]
  
  prop_rsaSignAndVerify :: Property
  prop_rsaSignAndVerify = property $ do
    msg <- forAll $ Gen.bytes (Range.linear 0 100)
    keylen <- forAll $ Gen.element ((`div` 8) <$> [2048, 3072, 4096])
    k <- evalIO $ genJWK (RSAGenParam keylen)
    alg_ <- forAll $ Gen.element [RS256, RS384, RS512, PS256, PS384, PS512]
    collect alg_
    msg' <- evalExceptT $ unwrapJOSE
      ( signJWS msg [(newJWSHeader (Protected, alg_), k)]
        >>= verifyJWS defaultValidationSettings k
        :: JOSE Error (PropertyT IO) B.ByteString
      )
    msg' === msg
  
  
  genCrv :: Gen Crv
  genCrv = Gen.element [P_256, P_384, P_521, Secp256k1]
  
  genOKPCrv :: Gen OKPCrv
  genOKPCrv = Gen.element [Ed25519, Ed448, X25519, X448]
  
  genKeyMaterialGenParam :: Gen KeyMaterialGenParam
  genKeyMaterialGenParam = Gen.choice
    [ ECGenParam <$> genCrv
    , RSAGenParam <$> Gen.element ((`div` 8) <$> [2048, 3072, 4096])
    , OctGenParam <$> liftA2 (+) (Gen.integral (Range.exponential 0 64)) (Gen.element [32, 48, 64])
    , OKPGenParam <$> genOKPCrv
    ]
  
  prop_bestJWSAlg :: Property
  prop_bestJWSAlg = property $ do
    msg <- forAll $ Gen.bytes (Range.linear 0 100)
  
    genParam <- forAll $ genKeyMaterialGenParam
    k <- evalIO $ genJWK genParam
  
    case bestJWSAlg k of
      Left (KeyMismatch _) -> discard   -- skip non-signing keys
      Left _ -> assert False
      Right alg_ -> do
        collect alg_
        msg' <- evalExceptT $ unwrapJOSE
          ( signJWS msg [(newJWSHeader (Protected, alg_), k)]
            >>= verifyJWS defaultValidationSettings k
            :: JOSE Error (PropertyT IO) B.ByteString
          )
        msg' === msg
  
  
  
  
  genRSAPrivateKeyOthElem :: Gen RSAPrivateKeyOthElem
  genRSAPrivateKeyOthElem =
    RSAPrivateKeyOthElem <$> genBase64Integer <*> genBase64Integer <*> genBase64Integer
  
  genRSAPrivateKeyOptionalParameters :: Gen RSAPrivateKeyOptionalParameters
  genRSAPrivateKeyOptionalParameters =
    RSAPrivateKeyOptionalParameters
      <$> genBase64Integer
      <*> genBase64Integer
      <*> genBase64Integer
      <*> genBase64Integer
      <*> genBase64Integer
      <*> Gen.maybe (Gen.nonEmpty (Range.linear 1 3) genRSAPrivateKeyOthElem)
  
  genRSAPrivateKeyParameters :: Gen RSAPrivateKeyParameters
  genRSAPrivateKeyParameters =
    RSAPrivateKeyParameters
      <$> genBase64Integer
      <*> Gen.maybe (genRSAPrivateKeyOptionalParameters)
  
  genRSAKeyParameters :: Gen RSAKeyParameters
  genRSAKeyParameters =
    RSAKeyParameters
      <$> genBase64Integer
      <*> genBase64Integer
      <*> Gen.maybe (genRSAPrivateKeyParameters)
  
  genDRG :: Gen ChaChaDRG
  genDRG = do
    let word64 = Gen.word64 Range.constantBounded
    seed <- (,,,,) <$> word64 <*> word64 <*> word64 <*> word64 <*> word64
    pure $ drgNewTest seed
  
  genECKeyParameters :: Gen ECKeyParameters
  genECKeyParameters = do
      drg <- genDRG
      crv <- genCrv
      let (k, _) = withDRG drg (genEC crv)
      includePrivate <- Gen.bool
      pure $ if includePrivate
        then k
        else (let Just a = view asPublicKey k in a)
  
  genOctKeyParameters :: Gen OctKeyParameters
  genOctKeyParameters = OctKeyParameters . Base64Octets <$> Gen.bytes (Range.linear 16 128)
  
  genOKPKeyParameters :: Gen OKPKeyParameters
  genOKPKeyParameters = do
    drg <- genDRG
    crv <- genOKPCrv
    let (k, _) = withDRG drg (genOKP crv)
    includePrivate <- Gen.bool
    pure $ if includePrivate
      then k
      else (let Just a = view asPublicKey k in a)
  
  genKeyMaterial' :: Gen KeyMaterial
  genKeyMaterial' = Gen.choice
    [ ECKeyMaterial <$> genECKeyParameters
    , RSAKeyMaterial <$> genRSAKeyParameters
    , OctKeyMaterial <$> genOctKeyParameters
    , OKPKeyMaterial <$> genOKPKeyParameters
    ]
  
  genBase64SHA1 :: Gen Base64SHA1
  genBase64SHA1 = Base64SHA1 <$> Gen.bytes (Range.singleton 20)
  
  genBase64SHA256 :: Gen Base64SHA256
  genBase64SHA256 = Base64SHA256 <$> Gen.bytes (Range.singleton 32)
  
  genJWK' :: Gen JWK
  genJWK' = do
    key <- genKeyMaterial'
    kid_ <- Gen.text (Range.linear 8 16) Gen.hexit
    x5t_ <- genBase64SHA1
    x5tS256_ <- genBase64SHA256
    pure $ fromKeyMaterial key
      & set jwkKid (Just kid_)
      & set jwkX5t (Just x5t_)
      & set jwkX5tS256 (Just x5tS256_)
  