module Nuthatch.Test.Oidc
  ( tests
  ) where

import           Control.Lens ((&), (?~))
import           Crypto.JOSE (runJOSE)
import           Crypto.JOSE.Compact (encodeCompact)
import           Crypto.JOSE.JWA.JWS (Alg (RS256, RS384))
import           Crypto.JOSE.JWK
  ( JWK
  , JWKSet (..)
  , KeyMaterialGenParam (RSAGenParam)
  , genJWK
  )
import           Crypto.JOSE.JWS (newJWSHeaderProtected)
import           Crypto.JWT
  ( Audience (..)
  , ClaimsSet
  , HasClaimsSet (..)
  , JWTError
  , NumericDate (..)
  , emptyClaimsSet
  , signJWT
  )
import qualified Data.Aeson as A
import qualified Data.Aeson.KeyMap as KM
import qualified Data.ByteString.Lazy as BL
import           Data.Either (isLeft, isRight)
import           Data.String (fromString)
import           Data.Text (Text)
import qualified Data.Text as T
import qualified Data.Text.Encoding as TE
import           Data.Time.Clock (UTCTime, addUTCTime, getCurrentTime)
import           Nuthatch.Config (Config (..), OidcConfig (..))
import           Nuthatch.Oidc (OidcIdentity (..), verifyIdToken)
import           Nuthatch.Test.Config (testConfig)
import           Test.Tasty (TestTree, testGroup)
import           Test.Tasty.HUnit ((@?=), assertBool, testCase)

tests :: TestTree
tests = testGroup "OIDC validation"
  [ testCase "accepts an exact RS256 ID token" testValidIdToken
  , testCase "accepts an exact configured audience set" testConfiguredAudiences
  , testCase "rejects invalid nonce, algorithm, issuer, signature, and audience"
      testIdTokenCryptographicRejections
  , testCase "rejects malformed, missing-time, future, expired, and overlong tokens"
      testIdTokenTimeAndSerializationRejections
  ]

oidcConfig :: OidcConfig
oidcConfig = cfgOidc testConfig

data TestIdTokenClaims = TestIdTokenClaims
  { testIdTokenRegistered      :: ClaimsSet
  , testIdTokenNonce           :: Text
  , testIdTokenAuthorizedParty :: Maybe Text
  , testIdTokenName            :: Maybe Text
  }
  deriving stock (Eq, Show)

instance HasClaimsSet TestIdTokenClaims where
  claimsSet update claims =
    fmap
      (\registered -> claims { testIdTokenRegistered = registered })
      (update (testIdTokenRegistered claims))

instance A.FromJSON TestIdTokenClaims where
  parseJSON value = A.withObject "test ID token" (\object ->
    TestIdTokenClaims
      <$> A.parseJSON value
      <*> object A..: "nonce"
      <*> object A..:? "azp"
      <*> object A..:? "name") value

instance A.ToJSON TestIdTokenClaims where
  toJSON claims =
    case A.toJSON (testIdTokenRegistered claims) of
      A.Object object ->
        let withName =
              maybe
                object
                (\name -> KM.insert "name" (A.String name) object)
                (testIdTokenName claims)
            withAuthorizedParty =
              maybe
                withName
                (\authorizedParty ->
                  KM.insert "azp" (A.String authorizedParty) withName)
                (testIdTokenAuthorizedParty claims)
        in A.Object $
          KM.insert
            "nonce"
            (A.String (testIdTokenNonce claims))
            withAuthorizedParty
      value -> value

testValidIdToken :: IO ()
testValidIdToken = do
  key <- genJWK (RSAGenParam 256)
  now <- getCurrentTime
  compact <- signTestIdToken key RS256
    (testIdClaims now "https://issuer.example" ["nuthatch"] "nonce")
  result <- verifyIdToken
    oidcConfig
    (JWKSet [key])
    now
    "nonce"
    compact
  result @?= Right OidcIdentity
    { identityIssuer = "https://issuer.example"
    , identitySubject = "upstream-subject"
    , identityDisplay = Just "Example User"
    }

testConfiguredAudiences :: IO ()
testConfiguredAudiences = do
  key <- genJWK (RSAGenParam 256)
  now <- getCurrentTime
  compact <- signTestIdToken key RS256
    ((testIdClaims now "https://issuer.example" ["project", "nuthatch"] "nonce")
      { testIdTokenAuthorizedParty = Just "nuthatch" })
  missingAuthorizedParty <- signTestIdToken key RS256
    (testIdClaims now "https://issuer.example" ["project", "nuthatch"] "nonce")
  result <- verifyIdToken
    oidcConfig { cfgOidcAdditionalAudiences = ["project"] }
    (JWKSet [key])
    now
    "nonce"
    compact
  rejected <- verifyIdToken
    oidcConfig { cfgOidcAdditionalAudiences = ["project"] }
    (JWKSet [key])
    now
    "nonce"
    missingAuthorizedParty
  assertBool "the configured audience set is accepted" (isRight result)
  assertBool "a multi-audience token without azp is rejected" (isLeft rejected)

testIdTokenCryptographicRejections :: IO ()
testIdTokenCryptographicRejections = do
  key <- genJWK (RSAGenParam 256)
  otherKey <- genJWK (RSAGenParam 256)
  now <- getCurrentTime
  valid <- signTestIdToken key RS256
    (testIdClaims now "https://issuer.example" ["nuthatch"] "nonce")
  wrongAlgorithm <- signTestIdToken key RS384
    (testIdClaims now "https://issuer.example" ["nuthatch"] "nonce")
  wrongIssuer <- signTestIdToken key RS256
    (testIdClaims now "https://other.example" ["nuthatch"] "nonce")
  wrongSignature <- signTestIdToken otherKey RS256
    (testIdClaims now "https://issuer.example" ["nuthatch"] "nonce")
  wrongAudience <- signTestIdToken key RS256
    (testIdClaims now "https://issuer.example" ["other"] "nonce")
  multipleAudience <- signTestIdToken key RS256
    (testIdClaims now "https://issuer.example" ["nuthatch", "other"] "nonce")
  let rejected =
        [ verifyIdToken oidcConfig (JWKSet [key]) now "other-nonce" valid
        , verifyIdToken oidcConfig (JWKSet [key]) now "nonce" wrongAlgorithm
        , verifyIdToken oidcConfig (JWKSet [key]) now "nonce" wrongIssuer
        , verifyIdToken oidcConfig (JWKSet [key]) now "nonce" wrongSignature
        , verifyIdToken oidcConfig (JWKSet [key]) now "nonce" wrongAudience
        , verifyIdToken oidcConfig (JWKSet [key]) now "nonce" multipleAudience
        ]
  results <- sequence rejected
  assertBool "every invalid cryptographic or binding case is rejected"
    (all isLeft results)

testIdTokenTimeAndSerializationRejections :: IO ()
testIdTokenTimeAndSerializationRejections = do
  key <- genJWK (RSAGenParam 256)
  now <- getCurrentTime
  missingIssuedAt <- signTestIdToken key RS256
    ((testIdClaims now "https://issuer.example" ["nuthatch"] "nonce")
      { testIdTokenRegistered =
          emptyClaimsSet
            & claimIss ?~ "https://issuer.example"
            & claimSub ?~ "upstream-subject"
            & claimAud ?~ Audience ["nuthatch"]
            & claimExp ?~ NumericDate (addUTCTime 300 now)
      })
  future <- signTestIdToken key RS256
    (testIdClaimsAt
      (addUTCTime 60 now)
      (addUTCTime 300 now)
      "https://issuer.example"
      ["nuthatch"]
      "nonce")
  expired <- signTestIdToken key RS256
    (testIdClaimsAt
      (addUTCTime (-600) now)
      (addUTCTime (-300) now)
      "https://issuer.example"
      ["nuthatch"]
      "nonce")
  overlong <- signTestIdToken key RS256
    (testIdClaimsAt
      now
      (addUTCTime 7200 now)
      "https://issuer.example"
      ["nuthatch"]
      "nonce")
  results <- sequence
    [ verifyIdToken oidcConfig (JWKSet [key]) now "nonce" "not.compact"
    , verifyIdToken oidcConfig (JWKSet [key]) now "nonce" missingIssuedAt
    , verifyIdToken oidcConfig (JWKSet [key]) now "nonce" future
    , verifyIdToken oidcConfig (JWKSet [key]) now "nonce" expired
    , verifyIdToken oidcConfig (JWKSet [key]) now "nonce" overlong
    ]
  assertBool "every invalid serialization or time case is rejected"
    (all isLeft results)

testIdClaims
  :: UTCTime
  -> Text
  -> [Text]
  -> Text
  -> TestIdTokenClaims
testIdClaims now issuer audiences nonce =
  testIdClaimsAt
    now
    (addUTCTime 300 now)
    issuer
    audiences
    nonce

testIdClaimsAt
  :: UTCTime
  -> UTCTime
  -> Text
  -> [Text]
  -> Text
  -> TestIdTokenClaims
testIdClaimsAt issuedAt expiresAt issuer audiences nonce =
  TestIdTokenClaims
    { testIdTokenRegistered =
        emptyClaimsSet
          & claimIss ?~ fromString (T.unpack issuer)
          & claimSub ?~ "upstream-subject"
          & claimAud ?~ Audience (map (fromString . T.unpack) audiences)
          & claimIat ?~ NumericDate issuedAt
          & claimExp ?~ NumericDate expiresAt
    , testIdTokenNonce = nonce
    , testIdTokenAuthorizedParty = Nothing
    , testIdTokenName = Just "Example User"
    }

signTestIdToken
  :: JWK
  -> Alg
  -> TestIdTokenClaims
  -> IO Text
signTestIdToken key algorithm claims = do
  result <- runJOSE $
    encodeCompact <$> signJWT key (newJWSHeaderProtected algorithm) claims
  case result of
    Left (err :: JWTError) -> fail (show err)
    Right compact -> pure (TE.decodeUtf8 (BL.toStrict compact))