module Main (main) where

import           Data.Aeson (Value (..), eitherDecodeStrict')
import qualified Data.Aeson.Key as Key
import qualified Data.Aeson.KeyMap as KeyMap
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as B8
import qualified Data.Text as Text
import           Nuthatch.Probe.Assert
import           Nuthatch.Probe.Http
import           Nuthatch.Probe.Oidc
import           Nuthatch.Probe.ProductClient
import           System.Environment (getArgs)

main :: IO ()
main = do
  context <- newProbeContext
  user <- requireEnv "NUTHATCH_OIDC_USER_A"
  session <- login context False user "product"
  getArgs >>= \case
    ["verify"] -> verifyProduct context session
    ["timeout"] -> verifyTimeout context session
    _ -> fail "expected verify or timeout mode"

verifyTimeout :: ProbeContext -> Session -> IO ()
verifyTimeout context session = do
  response <- authenticatedGetWith context session
    defaultRequestOptions
      { optionQuery =
          [ ("root", Just "personal")
          , ("q", Just "timeout-probe")
          ]
      }
    (probeBase context <> "/api/browse")
  assertEqual "timed search status" 200 (probeStatus response)
  assertContains "timed search bound" "\"partial\":true" (probeBody response)
  putStrLn "Timed search bound passed"

verifyProduct :: ProbeContext -> Session -> IO ()
verifyProduct context session = do
  verifyApplicationAssets context session
  verifyPaginationAndSearch context session
  verifyCanonicalPaths context session
  verifySafeDownload context session
  verifyRecursiveDelete context session
  putStrLn "Product behavior passed"

verifyApplicationAssets :: ProbeContext -> Session -> IO ()
verifyApplicationAssets context session = do
  sessionResponse <- authenticatedGet context session
    (probeBase context <> "/api/session")
  assertEqual "session endpoint status" 200 (probeStatus sessionResponse)
  assertContains "session user" "\"user_id\":" (probeBody sessionResponse)
  assertContains "session CSRF"
    ("\"csrf\":\"" <> sessionCsrf session <> "\"")
    (probeBody sessionResponse)
  application <- authenticatedGet context session (probeBase context <> "/app")
  assertEqual "application status" 200 (probeStatus application)
  csp <- maybe (fail "application omitted CSP") pure $
    headerValue "content-security-policy" application
  mapM_
    (\directive ->
      assertContains ("application CSP omitted " <> B8.unpack directive)
        directive
        csp)
    ["default-src 'none'", "script-src 'self'", "frame-ancestors 'none'"]
  asset <- between "src=\"/assets/" "\"" (B8.unpack (probeBody application))
  javascript <- authenticatedGet context session
    (probeBase context <> "/assets/" <> asset)
  assertEqual "asset status" 200 (probeStatus javascript)
  assertContains "asset content type" "javascript" $
    maybe BS.empty id (headerValue "content-type" javascript)
  assertEqual "asset nosniff"
    (Just "nosniff")
    (headerValue "x-content-type-options" javascript)
  let visible = probeBody application <> probeBody javascript
  mapM_
    (\value ->
      assert ("browser artifact exposed " <> value) $
        not (B8.pack value `BS.isInfixOf` visible))
    [ "NUTHATCH_S3"
    , "AWS_ACCESS_KEY"
    , "AWS_SECRET_ACCESS_KEY"
    , "client_secret"
    ]
  mapM_
    (\marker ->
      assertContains ("asset omitted " <> marker)
        (B8.pack marker)
        (probeBody javascript))
    [ "Drop a file here"
    , "Shared root"
    , "Retry same upload"
    , "Retry same delete"
    , "nuthatch-recursive-delete"
    , "URLSearchParams"
    ]
  unauthenticatedApp <- get (probeClient context) (probeBase context <> "/app")
  assertEqual "unauthenticated application" 303 (probeStatus unauthenticatedApp)
  unauthenticatedAsset <- get
    (probeClient context)
    (probeBase context <> "/assets/" <> asset)
  assertEqual "unauthenticated asset" 303 (probeStatus unauthenticatedAsset)

verifyPaginationAndSearch :: ProbeContext -> Session -> IO ()
verifyPaginationAndSearch context session = do
  uploadSucceeded context session "product-browse/Alpha.txt"
    "81000000-0000-4000-8000-000000000001" "alpha"
  replay <- upload context session "personal" "product-browse/Alpha.txt"
    "81000000-0000-4000-8000-000000000001" "alpha"
  requireSucceeded "product upload replay" replay
  uploadSucceeded context session "product-browse/beta.txt"
    "81000000-0000-4000-8000-000000000002" "beta"
  uploadSucceeded context session "product-browse/gamma.txt"
    "81000000-0000-4000-8000-000000000003" "gamma"
  firstPage <- authenticatedGetWith context session
    defaultRequestOptions
      { optionQuery =
          [ ("root", Just "personal")
          , ("prefix", Just "product-browse")
          ]
      }
    (probeBase context <> "/api/browse")
  requireStatus "first browse page" 200 firstPage
  assertContains "first page Alpha"
    "\"path\":\"product-browse/Alpha.txt\""
    (probeBody firstPage)
  assertContains "first page beta"
    "\"path\":\"product-browse/beta.txt\""
    (probeBody firstPage)
  assert "first page exceeded configured size" $
    not ("product-browse/gamma.txt" `BS.isInfixOf` probeBody firstPage)
  continuation <- stringField "continuation" (probeBody firstPage)
  secondPage <- authenticatedGetWith context session
    defaultRequestOptions
      { optionQuery =
          [ ("root", Just "personal")
          , ("prefix", Just "product-browse")
          , ("continuation", Just (B8.pack continuation))
          ]
      }
    (probeBase context <> "/api/browse")
  assertContains "second browse page"
    "\"path\":\"product-browse/gamma.txt\""
    (probeBody secondPage)
  mapM_
    (\(path, operationId, body) ->
      uploadSucceeded context session path operationId body)
    [ ("000-Needle-A.txt", "81000000-0000-4000-8000-000000000004", "one")
    , ("001-needle-b.txt", "81000000-0000-4000-8000-000000000005", "two")
    , ("002-filler.txt", "81000000-0000-4000-8000-000000000006", "three")
    , ("003-needle-after-bound.txt", "81000000-0000-4000-8000-000000000007", "four")
    ]
  searched <- authenticatedGetWith context session
    defaultRequestOptions
      { optionQuery =
          [ ("root", Just "personal")
          , ("q", Just "NEEDLE")
          ]
      }
    (probeBase context <> "/api/browse")
  mapM_
    (\value -> assertContains "bounded case-insensitive search" value (probeBody searched))
    [ "\"path\":\"000-Needle-A.txt\""
    , "\"path\":\"001-needle-b.txt\""
    , "\"partial\":true"
    ]
  assert "search crossed the configured scan bound" $
    not ("003-needle-after-bound.txt" `BS.isInfixOf` probeBody searched)

verifyCanonicalPaths :: ProbeContext -> Session -> IO ()
verifyCanonicalPaths context session = do
  mapM_
    (\url -> authenticatedGet context session url >>=
      requireStatus "non-canonical path" 400)
    [ probeBase context <> "/api/browse?root=Personal"
    , probeBase context <> "/api/browse?root=personal&prefix=a//b"
    , probeBase context <> "/api/download?root=personal&path=./file"
    ]
  decomposed <- authenticatedGet context session
    (probeBase context
      <> "/api/download?root=personal&path=cafe%CC%81.txt"
    )
  requireStatus "non-NFC path" 400 decomposed
  uploadSucceeded context session "encoded%2Fname.txt"
    "81000000-0000-4000-8000-000000000008" "encoded"
  encoded <- authenticatedGet context session
    (probeBase context
      <> "/api/download?root=personal&path=encoded%252Fname.txt"
    )
  assertEqual "encoded path body" "encoded" (probeBody encoded)
  uploadSucceeded context session "café.txt"
    "81000000-0000-4000-8000-000000000014" "unicode"
  unicode <- download context session "personal" "café.txt"
  assertEqual "Unicode path body" "unicode" (probeBody unicode)
  assertContains "Unicode filename"
    "filename*=UTF-8''caf%C3%A9.txt"
    (maybe BS.empty id (headerValue "content-disposition" unicode))

verifySafeDownload :: ProbeContext -> Session -> IO ()
verifySafeDownload context session = do
  let svg = "<svg onload=\"alert(1)\"></svg>"
  uploadSucceeded context session "active.svg"
    "81000000-0000-4000-8000-000000000009" svg
  ranged <- authenticatedGetWith context session
    defaultRequestOptions { optionHeaders = [("Range", "bytes=0-1")] }
    (probeBase context <> "/api/download?root=personal&path=active.svg")
  requireStatus "unsupported range" 416 ranged
  assertContains "range error"
    "\"error\":\"range_not_supported\""
    (probeBody ranged)
  downloaded <- download context session "personal" "active.svg"
  requireStatus "safe download" 200 downloaded
  let disposition =
        maybe BS.empty id (headerValue "content-disposition" downloaded)
  assertContains "download attachment" "attachment;" disposition
  assertContains "download filename" "filename*=UTF-8''active.svg" disposition
  assertEqual "download nosniff"
    (Just "nosniff")
    (headerValue "x-content-type-options" downloaded)
  assertContains "download sandbox"
    "sandbox"
    (maybe BS.empty id (headerValue "content-security-policy" downloaded))
  assertEqual "download body" svg (probeBody downloaded)

verifyRecursiveDelete :: ProbeContext -> Session -> IO ()
verifyRecursiveDelete context session = do
  mapM_
    (\(path, operationId, body) ->
      uploadSucceeded context session path operationId body)
    [ ("delete-tree/a.txt", "81000000-0000-4000-8000-000000000010", "a")
    , ("delete-tree/b.txt", "81000000-0000-4000-8000-000000000011", "b")
    , ("delete-tree/c.txt", "81000000-0000-4000-8000-000000000012", "c")
    ]
  planned <- deletePlan context session
  assertContains "first delete plan"
    "\"paths\":[\"delete-tree/a.txt\",\"delete-tree/b.txt\"]"
    (probeBody planned)
  assert "first delete plan omitted continuation" $
    not ("\"continuation\":null" `BS.isInfixOf` probeBody planned)
  let firstBatch =
        "{\"root\":\"personal\",\"csrf\":\""
        <> B8.unpack (sessionCsrf session)
        <> "\",\"items\":["
        <> "{\"path\":\"delete-tree/a.txt\",\"operation_id\":"
        <> "\"82000000-0000-4000-8000-000000000001\"},"
        <> "{\"path\":\"delete-tree/b.txt\",\"operation_id\":"
        <> "\"82000000-0000-4000-8000-000000000002\"}]}"
  first <- deleteBatch context session firstBatch
  requireBatchSucceeded "first delete batch" first
  replay <- deleteBatch context session firstBatch
  requireBatchSucceeded "delete batch replay" replay
  uploadSucceeded context session "delete-tree/d-concurrent.txt"
    "81000000-0000-4000-8000-000000000013" "d"
  restarted <- deletePlan context session
  assertContains "restarted delete plan"
    "\"paths\":[\"delete-tree/c.txt\",\"delete-tree/d-concurrent.txt\"]"
    (probeBody restarted)
  let secondBatch =
        "{\"root\":\"personal\",\"csrf\":\""
        <> B8.unpack (sessionCsrf session)
        <> "\",\"items\":["
        <> "{\"path\":\"delete-tree/c.txt\",\"operation_id\":"
        <> "\"82000000-0000-4000-8000-000000000003\"},"
        <> "{\"path\":\"delete-tree/d-concurrent.txt\",\"operation_id\":"
        <> "\"82000000-0000-4000-8000-000000000004\"}]}"
  second <- deleteBatch context session secondBatch
  requireBatchSucceeded "second delete batch" second
  empty <- deletePlan context session
  assertContains "empty delete plan" "\"paths\":[]" (probeBody empty)
  assertContains "empty delete continuation"
    "\"continuation\":null"
    (probeBody empty)

deletePlan :: ProbeContext -> Session -> IO ProbeResponse
deletePlan context session =
  authenticatedGetWith context session
    defaultRequestOptions
      { optionQuery =
          [ ("root", Just "personal")
          , ("path", Just "delete-tree")
          ]
      }
    (probeBase context <> "/api/delete-plan")

deleteBatch
  :: ProbeContext
  -> Session
  -> String
  -> IO ProbeResponse
deleteBatch context session body =
  postWith
    (probeClient context)
    defaultRequestOptions
      { optionCookies = Just (sessionCookies session)
      , optionHeaders = [("Content-Type", "application/json")]
      , optionBody = ByteBody (B8.pack body)
      }
    (probeBase context <> "/api/delete-batch")

requireBatchSucceeded :: String -> ProbeResponse -> IO ()
requireBatchSucceeded label response = do
  requireStatus label 200 response
  assertEqual label 2 $
    countOccurrences "\"result\":\"succeeded\"" (probeBody response)

uploadSucceeded
  :: ProbeContext
  -> Session
  -> String
  -> String
  -> BS.ByteString
  -> IO ()
uploadSucceeded context session path operationId body =
  upload context session "personal" path operationId body >>=
    requireSucceeded ("upload " <> path)

requireSucceeded :: String -> ProbeResponse -> IO ()
requireSucceeded label response = do
  requireStatus label 200 response
  assertContains label "\"result\":\"succeeded\"" (probeBody response)

requireStatus :: String -> Int -> ProbeResponse -> IO ()
requireStatus label expected response =
  assertEqual (label <> " status") expected (probeStatus response)

authenticatedGet
  :: ProbeContext
  -> Session
  -> String
  -> IO ProbeResponse
authenticatedGet context session =
  authenticatedGetWith context session defaultRequestOptions

authenticatedGetWith
  :: ProbeContext
  -> Session
  -> RequestOptions
  -> String
  -> IO ProbeResponse
authenticatedGetWith context session options =
  getWith
    (probeClient context)
    options { optionCookies = Just (sessionCookies session) }

stringField :: String -> BS.ByteString -> IO String
stringField name bytes =
  case eitherDecodeStrict' bytes of
    Right (Object object) ->
      case KeyMap.lookup (Key.fromString name) object of
        Just (String value) -> pure (Text.unpack value)
        _ -> fail ("JSON omitted string field " <> name)
    Right _ -> fail "expected JSON object"
    Left reason -> fail ("invalid JSON: " <> reason)

between :: String -> String -> String -> IO String
between start end value =
  case dropUntil start value of
    Nothing -> fail ("value omitted " <> start)
    Just afterStart ->
      case takeUntil end afterStart of
        Nothing -> fail ("value omitted " <> end)
        Just result -> pure result

dropUntil :: String -> String -> Maybe String
dropUntil needle value
  | needle `prefixOf` value = Just (drop (length needle) value)
dropUntil _ [] = Nothing
dropUntil needle (_ : rest) = dropUntil needle rest

takeUntil :: String -> String -> Maybe String
takeUntil needle = go []
  where
    go prefix value
      | needle `prefixOf` value = Just (reverse prefix)
    go _ [] = Nothing
    go prefix (next : rest) = go (next : prefix) rest

prefixOf :: String -> String -> Bool
prefixOf prefix value =
  take (length prefix) value == prefix

countOccurrences :: BS.ByteString -> BS.ByteString -> Int
countOccurrences needle input
  | BS.null input = 0
  | needle `BS.isPrefixOf` input =
      1 + countOccurrences needle (BS.drop (BS.length needle) input)
  | otherwise = countOccurrences needle (BS.drop 1 input)