module Nuthatch.Probe.Http ( CookieSession , HttpException , ProbeClient , ProbeResponse (..) , RequestBody (..) , RequestOptions (..) , cookieHeader , cookieValue , defaultRequestOptions , downloadTo , get , getWith , headWith , headerValue , headerValues , isConnectionFailure , newCookieSession , newProbeClient , postWith , putWith ) where import Control.Concurrent.MVar import qualified Data.ByteString as BS import qualified Data.ByteString.Char8 as B8 import qualified Data.ByteString.Lazy as BL import Data.List (find) import qualified Data.CaseInsensitive as CI import Network.Connection (TLSSettings (..)) import Network.HTTP.Client ( Cookie (..) , CookieJar , HttpException (..) , HttpExceptionContent (..) , Manager , Request , Response , brRead , applyBasicAuth , checkResponse , cookieJar , createCookieJar , destroyCookieJar , httpLbs , method , parseRequest , redirectCount , requestBody , requestHeaders , responseBody , responseCookieJar , responseHeaders , responseStatus , responseTimeout , responseTimeoutMicro , withResponse ) import qualified Network.HTTP.Client as HTTP import Network.HTTP.Client.TLS (mkManagerSettings) import Network.HTTP.Types ( Header , Method , hContentType , methodGet , methodHead , methodPost , methodPut , statusCode ) import Network.HTTP.Types.URI (renderSimpleQuery) import Network.TLS (defaultSupported) import System.IO (IOMode (WriteMode), withBinaryFile) newtype ProbeClient = ProbeClient Manager newtype CookieSession = CookieSession (MVar CookieJar) data ProbeResponse = ProbeResponse { probeStatus :: Int , probeHeaders :: [Header] , probeBody :: BS.ByteString } deriving stock (Eq, Show) data RequestBody = NoBody | ByteBody BS.ByteString | FileBody FilePath | FormBody [(BS.ByteString, BS.ByteString)] deriving stock (Eq, Show) data RequestOptions = RequestOptions { optionHeaders :: [Header] , optionBody :: RequestBody , optionCookies :: Maybe CookieSession , optionBasicAuth :: Maybe (BS.ByteString, BS.ByteString) , optionQuery :: [(BS.ByteString, Maybe BS.ByteString)] , optionTimeoutMicros :: Maybe Int } defaultRequestOptions :: RequestOptions defaultRequestOptions = RequestOptions { optionHeaders = [] , optionBody = NoBody , optionCookies = Nothing , optionBasicAuth = Nothing , optionQuery = [] , optionTimeoutMicros = Nothing } newProbeClient :: IO ProbeClient newProbeClient = ProbeClient <$> HTTP.newManager (mkManagerSettings TLSSettingsSimple { settingDisableCertificateValidation = True , settingDisableSession = False , settingUseServerName = True , settingClientSupported = defaultSupported } Nothing) newCookieSession :: IO CookieSession newCookieSession = CookieSession <$> newMVar (createCookieJar []) get :: ProbeClient -> String -> IO ProbeResponse get client = getWith client defaultRequestOptions getWith :: ProbeClient -> RequestOptions -> String -> IO ProbeResponse getWith client = perform client methodGet postWith :: ProbeClient -> RequestOptions -> String -> IO ProbeResponse postWith client = perform client methodPost putWith :: ProbeClient -> RequestOptions -> String -> IO ProbeResponse putWith client = perform client methodPut headWith :: ProbeClient -> RequestOptions -> String -> IO ProbeResponse headWith client = perform client methodHead downloadTo :: ProbeClient -> RequestOptions -> FilePath -> String -> IO ProbeResponse downloadTo (ProbeClient manager) options outputPath url = do request <- buildRequest methodGet options url withResponse request manager $ \response -> do updateCookies options response withBinaryFile outputPath WriteMode $ \handle -> drain handle (responseBody response) pure (responseWithoutBody response) where drain handle reader = do chunk <- brRead reader if BS.null chunk then pure () else BS.hPut handle chunk >> drain handle reader perform :: ProbeClient -> Method -> RequestOptions -> String -> IO ProbeResponse perform (ProbeClient manager) requestMethod options url = do request <- buildRequest requestMethod options url response <- httpLbs request manager updateCookies options response pure ProbeResponse { probeStatus = statusCode (responseStatus response) , probeHeaders = responseHeaders response , probeBody = BL.toStrict (responseBody response) } buildRequest :: Method -> RequestOptions -> String -> IO Request buildRequest requestMethod options url = do initial <- parseRequest url body <- toRequestBody (optionBody options) storedCookies <- traverse readCookies (optionCookies options) let formHeaders = case optionBody options of FormBody _ -> [(hContentType, "application/x-www-form-urlencoded")] _ -> [] queried = if null (optionQuery options) then initial else HTTP.setQueryString (optionQuery options) initial let configured = queried { method = requestMethod , requestHeaders = formHeaders <> optionHeaders options , requestBody = body , cookieJar = storedCookies , redirectCount = 0 , checkResponse = \_ _ -> pure () , responseTimeout = maybe (responseTimeout initial) responseTimeoutMicro (optionTimeoutMicros options) } pure $ maybe configured (\(username, password) -> applyBasicAuth username password configured) (optionBasicAuth options) toRequestBody :: RequestBody -> IO HTTP.RequestBody toRequestBody NoBody = pure (HTTP.RequestBodyBS BS.empty) toRequestBody (ByteBody body) = pure (HTTP.RequestBodyBS body) toRequestBody (FileBody path) = HTTP.streamFile path toRequestBody (FormBody fields) = pure . HTTP.RequestBodyBS $ renderSimpleQuery False fields updateCookies :: RequestOptions -> Response body -> IO () updateCookies options response = case optionCookies options of Nothing -> pure () Just (CookieSession storage) -> do _ <- swapMVar storage (responseCookieJar response) pure () readCookies :: CookieSession -> IO CookieJar readCookies (CookieSession storage) = readMVar storage cookieValue :: CookieSession -> BS.ByteString -> IO BS.ByteString cookieValue session expectedName = do cookies <- destroyCookieJar <$> readCookies session maybe (fail ("cookie not found: " <> B8.unpack expectedName)) (pure . cookie_value) (find ((== expectedName) . cookie_name) cookies) cookieHeader :: CookieSession -> IO BS.ByteString cookieHeader session = do cookies <- destroyCookieJar <$> readCookies session pure . B8.intercalate "; " $ [cookie_name cookie <> "=" <> cookie_value cookie | cookie <- cookies] headerValue :: BS.ByteString -> ProbeResponse -> Maybe BS.ByteString headerValue name response = snd <$> find ((== CI.mk name) . fst) (probeHeaders response) headerValues :: BS.ByteString -> ProbeResponse -> [BS.ByteString] headerValues name response = map snd (filter ((== CI.mk name) . fst) (probeHeaders response)) responseWithoutBody :: Response body -> ProbeResponse responseWithoutBody response = ProbeResponse { probeStatus = statusCode (responseStatus response) , probeHeaders = responseHeaders response , probeBody = BS.empty } isConnectionFailure :: HttpException -> Bool isConnectionFailure (HttpExceptionRequest _ (ConnectionFailure _)) = True isConnectionFailure _ = False