{-# LANGUAGE ApplicativeDo #-} {-# LANGUAGE CPP #-} {-# LANGUAGE DuplicateRecordFields #-} {-# LANGUAGE LambdaCase #-} {-# LANGUAGE MultiWayIf #-} {-# LANGUAGE NamedFieldPuns #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE PatternSynonyms #-} {-# LANGUAGE RecordWildCards #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE StrictData #-} {-# LANGUAGE TypeApplications #-} {-# LANGUAGE ViewPatterns #-} module Main where import Control.Exception import Data.Foldable import Data.List.NonEmpty (NonEmpty) import Data.Maybe import Data.Monoid import System.Exit import System.IO import qualified Data.List.NonEmpty as NE import qualified Data.Version as V import qualified System.Timeout as Timeout import Control.Lens hiding ((.=)) import Data.Time import Data.Time.Clock.POSIX import Network.HTTP.Client (HttpException) import Network.Socket import System.FilePath (()) import System.IO.Temp import System.Process import System.PosixCompat.User (getEffectiveUserName) import qualified Data.ByteString.Lazy as BL import qualified Data.ByteString.Lazy.Char8 as BL8 import qualified Data.Text as T import qualified Data.Text.Encoding.Error as TE import qualified Data.Text.Lazy as TL import qualified Data.Text.Lazy.Builder as TL import qualified Data.Text.Lazy.Encoding as TL import qualified Network.HTTP.Types.Status as HT import qualified Network.Socket.ByteString.Lazy as Socket import qualified Network.Wreq as W import qualified Options.Applicative as O import System.Process.Kill (killProcessTree) import qualified Paths_sensu_run as Paths import Data.Aeson #if MIN_VERSION_aeson(1, 2, 2) hiding (Options) #endif main :: IO () main = do issued <- getCurrentTime opts <- O.execParser $ O.info (O.helper <*> options) O.fullDesc case opts of ShowVersion -> do putStrLn $ "sensu-run " ++ V.showVersion Paths.version exitSuccess RunOptions {..} -> withSystemTempFile "sensu-run.XXX" $ \path hdl -> do executed <- getCurrentTime rawStatus <- try $ bracket (startProcess cmdspec hdl) (\ph -> do terminateProcess ph killProcessTree ph waitForProcess ph) (withTimeout timeout . waitForProcess) hClose hdl exited <- getCurrentTime rawOutput <- BL.readFile path user <- T.pack <$> getEffectiveUserName let encoded = encode CheckResult { command = cmdspec , output = TL.toLazyText $ mconcat [ TL.fromLazyText (TL.decodeUtf8With TE.lenientDecode rawOutput) , "\n" , TL.fromString $ case rawStatus of Left ioe -> show (ioe :: IOException) Right Nothing -> "sensu-run: timed out" Right _ -> mempty ] , status = case rawStatus of Right (Just ExitSuccess) -> OK Right (Just ExitFailure {}) -> CRITICAL _ -> UNKNOWN , duration = diffUTCTime exited executed , .. } if dryRun then BL8.putStrLn encoded else case endpoint of ClientSocketInput port -> sendToClientSocketInput port encoded SensuServer urls -> sendToSensuServer urls encoded case rawStatus of Left ioe -> do hPutStrLn stderr $ show ioe exitFailure Right (Just ExitSuccess) -> exitSuccess Right Nothing -> do hPutStrLn stderr $ showCmdSpec cmdspec ++ " timed out" exitFailure Right (Just ExitFailure {}) -> exitFailure sendToClientSocketInput :: PortNumber -- ^ Listening port of Sensu client socket -> BL8.ByteString -- ^ Payload -> IO () sendToClientSocketInput port payload = bracket open close $ \sock -> do localhost <- inet_addr "127.0.0.1" connect sock $ SockAddrInet port localhost Socket.sendAll sock payload `catch` \(ioe :: IOException) -> do hPutStrLn stderr $ "Failed to write results to localhost:" ++ show port ++ " (" ++ show ioe ++ ")" exitFailure where open = socket AF_INET Stream defaultProtocol sendToSensuServer :: NonEmpty String -- ^ Sensu server base URLs -> BL8.ByteString -- ^ Payload -> IO () sendToSensuServer urls payload = foldr go (handleError "no more retry") urls `catch` \(e :: HttpException) -> handleError (show e) where go url retry = do resp <- W.postWith params (url "results") payload let status = resp ^. W.responseStatus if | HT.statusIsClientError status -> handleError $ show status | HT.statusIsServerError status -> retry | HT.statusIsSuccessful status -> return () | otherwise -> fail $ "sendToSensuServer: unexpected status " ++ show status params = W.defaults & W.header "Content-Type" .~ ["application/json"] handleError reason = do hPutStrLn stderr $ "Failed to POST results to Sensu server (" ++ reason ++ ")" exitFailure startProcess :: CmdSpec -> Handle -> IO ProcessHandle startProcess cmdspec hdl = do (_, _, _, ph) <- createProcess CreateProcess { cmdspec , cwd = Nothing , env = Nothing , std_in = Inherit , std_out = UseHandle hdl , std_err = UseHandle hdl , close_fds = False , create_group = True -- necessary to not kill sensu-run itself , delegate_ctlc = False , detach_console = False , create_new_console = False , new_session = False , child_group = Nothing , child_user = Nothing #if MIN_VERSION_process(1, 5, 0) , use_process_jobs = True #endif } return ph withTimeout :: Maybe NominalDiffTime -> IO a -> IO (Maybe a) withTimeout time io = case time of Just n -> Timeout.timeout (seconds n) io Nothing -> Just <$> io where seconds n = round $ n * 10 ^ (6 :: Int) data Options = ShowVersion | RunOptions { name :: T.Text , cmdspec :: CmdSpec , source :: Maybe T.Text , ttl :: Maybe NominalDiffTime , timeout :: Maybe NominalDiffTime , handlers :: [T.Text] , endpoint :: Endpoint , dryRun :: Bool } data Endpoint = ClientSocketInput PortNumber -- ^ Local client socket input | SensuServer (NonEmpty String) -- ^ Sensu server API or a client HTTP socket -- -- Multiple HTTP endpoints can be specified. sensu-run retries sequentially -- until it succeeds. By default Sensu servers listen on port 4567 and -- client HTTP sockets listen on 3031. options :: O.Parser Options options = asum [ runOptions , ShowVersion <$ O.switch (O.long "version" <> O.short 'v') ] where runOptions = do name <- textOption $ mconcat [ O.short 'n' , O.long "name" , O.metavar "NAME" , O.help "The name of the check" ] source <- O.optional $ textOption $ mconcat [ O.long "source" , O.metavar "SOURCE" , O.help $ unlines [ "The check source, used to create a JIT Sensu client for an" , "external resource" ] ] ttl <- durationOption $ mconcat [ O.long "ttl" , O.metavar "SECONDS" , O.help "The time to live in seconds until check results are considered stale" ] timeout <- durationOption $ mconcat [ O.long "timeout" , O.metavar "SECONDS" , O.help "The check executaion duration timeout in seconds" ] handlers <- O.some $ textOption $ mconcat [ O.long "handler" , O.metavar "HANDLER" , O.help "Sensu event handler(s) to use for events created by the check" ] endpoint <- asum [ ClientSocketInput <$> portOption , SensuServer . NE.fromList <$> O.some serverOption ] dryRun <- O.switch $ mconcat [ O.long "dry-run" , O.long "dry" ] cmdspec <- cmdSpecOption return RunOptions {..} textOption m = T.pack <$> O.strOption m durationOption m = fmap (realToFrac @Double) <$> O.optional (O.option O.auto m) cmdSpecOption = cmdSpec <$> O.optional (O.switch $ mconcat [ O.short 's' , O.long "shell" , O.help "Execute the command using the shell" ]) <*> O.some (O.strArgument $ O.metavar "COMMAND") where cmdSpec (fromMaybe False -> isShell) args | isShell = ShellCommand (unwords args) | otherwise = RawCommand (head args) (tail args) portOption = O.option O.auto $ mconcat [ O.long "port" , O.metavar "PORT" , O.help "Send results to the local sensu-client listening on the specified port" , O.showDefault , O.value 3030 ] serverOption = O.strOption $ mconcat [ O.long "server" , O.metavar "URL" , O.help "Send results to the specified Sensu server" ] data CheckResult = CheckResult { name :: T.Text , command :: CmdSpec , status :: ExitCode , source :: Maybe T.Text , issued :: UTCTime , executed :: UTCTime , duration :: NominalDiffTime , output :: TL.Text , handlers :: [T.Text] , user :: T.Text } pattern OK :: ExitCode pattern OK = ExitSuccess pattern WARNING :: ExitCode pattern WARNING = ExitFailure 1 pattern CRITICAL :: ExitCode pattern CRITICAL = ExitFailure 2 pattern UNKNOWN :: ExitCode pattern UNKNOWN = ExitFailure 3 instance ToJSON CheckResult where toJSON = object . checkResultKeyValue toEncoding = pairs . mconcat . checkResultKeyValue checkResultKeyValue :: KeyValue a => CheckResult -> [a] checkResultKeyValue CheckResult {..} = addOptional "source" source [ "name" .= name , "command" .= showCmdSpec command , "issued" .= (floor (utcTimeToPOSIXSeconds issued) :: Int) , "executed" .= (floor (utcTimeToPOSIXSeconds executed) :: Int) , "duration" .= (realToFrac duration :: Double) , "status" .= statusToInt status , "output" .= output , "handlers" .= handlers , "user" .= user ] where addOptional key val ps = maybe ps (\val' -> key .= val' : ps) val statusToInt ExitSuccess = 0 statusToInt (ExitFailure n) = n showCmdSpec :: CmdSpec -> String showCmdSpec = \case ShellCommand cmd -> cmd RawCommand cmd args -> unwords $ cmd:args