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 (..))
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
  ]

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
    testConfig
    (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
    testConfig { cfgOidcAdditionalAudiences = ["project"] }
    (JWKSet [key])
    now
    "nonce"
    compact
  rejected <- verifyIdToken
    testConfig { 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 testConfig (JWKSet [key]) now "other-nonce" valid
        , verifyIdToken testConfig (JWKSet [key]) now "nonce" wrongAlgorithm
        , verifyIdToken testConfig (JWKSet [key]) now "nonce" wrongIssuer
        , verifyIdToken testConfig (JWKSet [key]) now "nonce" wrongSignature
        , verifyIdToken testConfig (JWKSet [key]) now "nonce" wrongAudience
        , verifyIdToken testConfig (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 testConfig (JWKSet [key]) now "nonce" "not.compact"
    , verifyIdToken testConfig (JWKSet [key]) now "nonce" missingIssuedAt
    , verifyIdToken testConfig (JWKSet [key]) now "nonce" future
    , verifyIdToken testConfig (JWKSet [key]) now "nonce" expired
    , verifyIdToken testConfig (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))