diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright (c) 2012, Jason Hickner + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of Jason Hickner nor the names of other + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/minions.cabal b/minions.cabal new file mode 100644 --- /dev/null +++ b/minions.cabal @@ -0,0 +1,29 @@ +-- Initial minions.cabal generated by cabal init. For further +-- documentation, see http://haskell.org/cabal/users-guide/ + +name: minions +version: 0.1.0.0 +synopsis: A fast parallel ssh tool +-- description: +homepage: http://github.com/jhickner/minions +license: BSD3 +license-file: LICENSE +author: Jason Hickner +maintainer: jhickner@gmail.com +-- copyright: +category: Network +build-type: Simple +cabal-version: >=1.8 + +source-repository head + type: git + location: git://github.com/jhickner/minions + +executable minions + main-is: minions.hs + GHC-options: -Wall -O2 -threaded + build-depends: base ==4.5.*, + ansi-terminal ==0.5.*, + process ==1.1.*, + time ==1.4.*, + MissingH ==1.2.* diff --git a/minions.hs b/minions.hs new file mode 100644 --- /dev/null +++ b/minions.hs @@ -0,0 +1,144 @@ +{-# LANGUAGE RecordWildCards #-} + +module Main where + +import System.Environment (getArgs) +import System.Exit (ExitCode(..), exitFailure, exitSuccess) +import System.Console.ANSI +import System.Console.GetOpt +import System.Process (readProcessWithExitCode) +import System.Timeout (timeout) +import System.IO (stderr, hPutStrLn) + +import Control.Concurrent +import Control.Monad + +import Data.Time.Clock +import Data.String.Utils +import Text.Printf + + +data Task = Task + { taskHost :: String + , taskCmd :: [String] + } deriving Show + +data Result = Result + { resHost :: String + , resPayload :: Either String String + , resTime :: NominalDiffTime + } deriving Show + +-- | time the given IO action (clock time) and return a tuple +-- of the execution time and the result +timeIO :: IO a -> IO (NominalDiffTime, a) +timeIO ioa = do + t1 <- getCurrentTime + a <- ioa + t2 <- getCurrentTime + return (diffUTCTime t2 t1, a) + +runTasks :: [Task] -> Int -> (Result -> IO a) -> IO () +runTasks ts tmout handler = do + out <- newChan + forM_ ts $ run out + replicateM_ (length ts) (readChan out >>= handler) + where run ch t = forkIO $ do + res <- runTask t tmout + writeChan ch res + +runTask :: Task -> Int -> IO Result +runTask Task{..} tmout = do + (time, res) <- timeIO . timeout tmout $ readProcessWithExitCode "ssh" args [] + case res of + Nothing -> output time $ Left "timed out\n" + Just (code, stdout', stderr') -> + output time $ case code of + ExitSuccess -> Right stdout' + ExitFailure _ -> Left stderr' + where args = "-T" : taskHost : taskCmd + output time payload = return $ Result taskHost payload time + +mkTasks :: [String] -> [String] -> [Task] +mkTasks hosts cmd = map f hosts + where f h = Task h cmd + + +---------------- +-- PRINTING +---------------- + +putColorLn :: ColorIntensity -> Color -> String -> IO () +putColorLn intensity color str = do + setSGR [SetColor Foreground intensity color] + putStrLn str + setSGR [] + +printResult :: Result -> IO () +printResult Result{..} = do + putColorLn Dull Blue (resHost ++ t) + case resPayload of + Left s -> putColorLn Vivid Red s + Right s -> putStrLn s + where t = printf " (%.1fs)" (realToFrac resTime :: Double) + +printShortResult :: Result -> IO () +printShortResult Result{..} = + case resPayload of + Left s -> putColorLn Vivid Red $ resHost ++ t ++ rstrip s + Right s -> putStrLn $ resHost ++ t ++ rstrip s + where t = printf " (%.1fs): " (realToFrac resTime :: Double) + + +---------------- +-- ARGS +---------------- + +data Options = Options + { oHelp :: Bool + , oHosts :: FilePath + , oTimeout :: Int + , oHandler :: Result -> IO () + } + +defaultOptions :: Options +defaultOptions = Options False "" (10 * 1000000) printResult + +options :: [OptDescr (Options -> Options)] +options = + [ Option [] ["help"] + (NoArg (\opts -> opts { oHelp = True })) + "display this help" + , Option "h" ["hosts"] + (ReqArg (\f opts -> opts { oHosts = f }) "FILE") + "FILE containing ssh host names (one per line)" + , Option "t" ["timeout"] + (ReqArg (\f opts -> opts { oTimeout = (read f :: Int) * 1000000 }) + "SECONDS") + "ssh timeout in SECONDS (default 10)" + , Option "s" ["short"] + (NoArg (\opts -> opts { oHandler = printShortResult })) + "display results in short format" + ] + +parseArgs :: IO (Options, [String]) +parseArgs = do + argv <- getArgs + case getOpt RequireOrder options argv of + (o, n, []) -> return (foldl (flip id) defaultOptions o, n) + (_, _, es) -> showError $ concat es + +showError :: String -> IO a +showError msg = hPutStrLn stderr (msg ++ header) >> exitFailure + +header :: String +header = usageInfo "Usage: minions [-hst] command" options + +main :: IO () +main = do + (Options{..}, cmd) <- parseArgs + when oHelp (putStrLn header >> exitSuccess) + when (oHosts == "") (showError "Please specify a hostname file with -h\n") + when (null cmd) (showError "Please specify a command to run\n") + hosts <- (filter ((not . null) . strip) . lines) `fmap` readFile oHosts + runTasks (mkTasks hosts cmd) oTimeout oHandler