nuthatch/test/Nuthatch/Probe/Curl.hs

module Nuthatch.Probe.Curl
  ( CurlResponse (..)
  , CurlTransportFailure (..)
  , cookieValue
  , curl
  , curlToFile
  , headerValue
  , headerValues
  ) where

import           Control.Exception (Exception, throwIO)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as B8
import           Data.Char (toLower)
import           Data.List (find)
import           Data.Maybe (mapMaybe)
import           System.Exit (ExitCode (..))
import           System.IO.Temp (withSystemTempDirectory)
import           System.Process (readProcessWithExitCode)
import           Text.Read (readMaybe)

data CurlResponse = CurlResponse
  { curlStatus  :: Int
  , curlHeaders :: [(BS.ByteString, BS.ByteString)]
  , curlBody    :: BS.ByteString
  } deriving stock (Eq, Show)

data CurlTransportFailure = CurlTransportFailure
  { curlTransportExitCode :: Int
  , curlTransportStderr   :: String
  } deriving stock (Eq, Show)

instance Exception CurlTransportFailure

curl :: [String] -> IO CurlResponse
curl arguments =
  withSystemTempDirectory "nuthatch-curl" $ \directory -> do
    let bodyPath = directory <> "/body"
    response <- runCurl bodyPath arguments
    body <- BS.readFile bodyPath
    pure response { curlBody = body }

curlToFile :: FilePath -> [String] -> IO CurlResponse
curlToFile outputPath arguments =
  runCurl outputPath arguments

runCurl :: FilePath -> [String] -> IO CurlResponse
runCurl outputPath arguments =
  withSystemTempDirectory "nuthatch-curl-headers" $ \directory -> do
    let headersPath = directory <> "/headers"
        common =
          [ "--insecure"
          , "--silent"
          , "--show-error"
          , "--dump-header", headersPath
          , "--output", outputPath
          , "--write-out", "%{http_code}"
          ]
    (exitCode, statusText, stderrText) <-
      readProcessWithExitCode "curl" (common <> arguments) ""
    status <- case readMaybe statusText of
      Just value -> pure value
      Nothing ->
        fail ("curl returned an invalid status: " <> statusText)
    case exitCode of
      ExitSuccess -> pure ()
      ExitFailure code ->
        throwIO CurlTransportFailure
          { curlTransportExitCode = code
          , curlTransportStderr = stderrText
          }
    headers <- parseHeaders <$> BS.readFile headersPath
    pure CurlResponse
      { curlStatus = status
      , curlHeaders = headers
      , curlBody = BS.empty
      }

parseHeaders :: BS.ByteString -> [(BS.ByteString, BS.ByteString)]
parseHeaders input =
  case filter (not . null) (splitBlocks (B8.lines input)) of
    [] -> []
    blocks -> mapMaybe parseHeader (last blocks)
  where
    splitBlocks = foldr splitLine [[]]
    splitLine line (block : blocks)
      | BS.null (stripCarriageReturn line) = [] : block : blocks
      | otherwise = (stripCarriageReturn line : block) : blocks
    splitLine _ [] = []

parseHeader :: BS.ByteString -> Maybe (BS.ByteString, BS.ByteString)
parseHeader line =
  case B8.break (== ':') line of
    (_, remainder) | BS.null remainder -> Nothing
    (name, remainder) ->
      Just (lower name, B8.dropWhile (== ' ') (BS.drop 1 remainder))

headerValue :: BS.ByteString -> CurlResponse -> Maybe BS.ByteString
headerValue name response =
  snd <$> find ((== lower name) . fst) (curlHeaders response)

headerValues :: BS.ByteString -> CurlResponse -> [BS.ByteString]
headerValues name response =
  map snd (filter ((== lower name) . fst) (curlHeaders response))

lower :: BS.ByteString -> BS.ByteString
lower = B8.map toLower

stripCarriageReturn :: BS.ByteString -> BS.ByteString
stripCarriageReturn value =
  case B8.unsnoc value of
    Just (prefix, '\r') -> prefix
    _ -> value

cookieValue :: FilePath -> BS.ByteString -> IO BS.ByteString
cookieValue path expectedName = do
  contents <- B8.lines <$> BS.readFile path
  case mapMaybe parseCookie contents of
    cookies ->
      maybe
        (fail ("cookie not found: " <> B8.unpack expectedName))
        pure
        (lookup expectedName cookies)

parseCookie :: BS.ByteString -> Maybe (BS.ByteString, BS.ByteString)
parseCookie line =
  let normalized =
        maybe line id (BS.stripPrefix "#HttpOnly_" line)
      fields = B8.split '\t' normalized
  in case drop 5 fields of
      name : value : _ -> Just (name, value)
      _ -> Nothing