ztail-1.0: Main.hs

{-# LANGUAGE CPP #-}
import System.Exit
import Control.Monad
import qualified System.Environment
import qualified System.Console.GetOpt as GetOpt
import System.Posix.Signals (installHandler, sigINT, Handler(..))
import qualified GHC.IOBase
import qualified Control.Concurrent
import qualified Control.Exception
import qualified Data.IORef
import qualified Text.Regex
#ifdef INOTIFY
import qualified System.INotify as INotify
#endif

import Util
import Display
import TailTypes
import Tail
import TailHandle

data Options = Options{
  optionTails :: [Tail],
  optionTail :: Tail,
  optionMatch :: TailMatch
}

defaultTail = Tail{
  tailTarg = undefined, -- read "-",
  tailPollInterval = read "5",
  tailReopenInterval = read "0",
#ifdef INOTIFY
  tailPollINotify = True,
  tailReopenINotify = False,
#endif
  tailBegin = False,
  tailTimeFmt = "%c",
  tailMatches = []
}

defaultOptions = Options{
  optionTails = [],
  optionTail = defaultTail,
  optionMatch = MatchAll
}

set_opt :: (Tail -> Tail) -> Options -> Options
set_opt p o = o{ optionTail = p $ optionTail o }

set_match :: TailMatch -> Options -> Options
set_match m o = o{ optionMatch = m }

add_action :: TailAction -> Options -> Options
add_action a o = set_opt add o where
  add t = t{ tailMatches = (optionMatch o, a) : (tailMatches t) }

prog_header = "Usage: ztail [OPTIONS] FILE ...\n\
Follow the specified files (ala tail -f).  FILE may be a path, '-' for stdin,\n\
or '&N' for file descriptor N.  OPTIONS apply only to the following FILE\n\
(except for -irt), and match options (-amn) apply to all following actions\n\
(-hcdse).  Actions involving TEXT can contain the following \\-escapes:\n\
    \\0 current file        \\@ current time (from -t)\n\
    \\_ current line        \\` \\' pre- and post-matching text\n\
    \\& matching text       \\N (1-9) group in match\n\
\&" --"
prog_usage = GetOpt.usageInfo prog_header prog_options

prog_options :: [GetOpt.OptDescr (Options -> Options)]
prog_options = [
    GetOpt.Option "b" ["begin"]
      (GetOpt.NoArg (set_opt $ \p -> p{ tailBegin = True }))
      ("start reading at the beginning of the file (rather than only new lines at the end)"),
    GetOpt.Option "i" ["interval"]
      (GetOpt.ReqArg (\i -> set_opt $ \p -> p{ tailPollInterval = read i }) "INT")
      ("poll for data every INT seconds [" ++ show (tailPollInterval defaultTail) ++ "] on all following files"),
    GetOpt.Option "r" ["reopen"]
      (GetOpt.OptArg (\i -> set_opt $ \p -> p{ tailReopenInterval = maybe (tailPollInterval p) read i }) "INT")
      ("check file name (like tail -F) every INT seconds or every poll"),
    GetOpt.Option "t" ["timefmt"]
      (GetOpt.ReqArg (\t -> set_opt $ \p -> p{ tailTimeFmt = t }) "FMT")
      ("set time format for \\@ substitution (in strftime(3)) [" ++ tailTimeFmt defaultTail ++ "]"),
    GetOpt.Option "T" ["timestamp"]
      (GetOpt.OptArg (maybe id $ \t -> add_action (ActionSubst "\\@ \\_") . set_opt (\p -> p{ tailTimeFmt = t })) "FMT")
      ("timestamp with FMT; equivalent to: [-t FMT] -h '\\@ '"),
#ifdef INOTIFY
    GetOpt.Option "I" ["inotify"]
      (GetOpt.NoArg (set_opt $ \p -> p{ tailPollINotify = True, tailPollInterval = read "0" }))
      ("use inotify to poll for new data (only for regular files)"),
    GetOpt.Option "R" ["ireopen"]
      (GetOpt.NoArg (set_opt $ \p -> p{ tailReopenINotify = True }))
      ("use inotify to monitor file renames (only for preexisting, leaf files)"),
#endif

    GetOpt.Option "a" ["all"]
      (GetOpt.NoArg (set_match MatchAll))
      ("perform following action for every line from this FILE (default)"),
    GetOpt.Option "m" ["match"]
      (GetOpt.ReqArg (\m -> set_match $ MatchRegex $ Text.Regex.mkRegexWithOpts m False True) "REGEX")
      ("perform following action for each line matching REGEX"),
    GetOpt.Option "M" ["imatch"]
      (GetOpt.ReqArg (\m -> set_match $ MatchRegex $ Text.Regex.mkRegexWithOpts m False False) "REGEX")
      ("perform following action for each line matching REGEX (case-insensitive)"),
    GetOpt.Option "n" ["no-match"]
      (GetOpt.ReqArg (\m -> set_match $ MatchNotRegex $ Text.Regex.mkRegexWithOpts m False True) "REGEX")
      ("perform following action for each line not matching REGEX"),
    GetOpt.Option "N" ["no-imatch"]
      (GetOpt.ReqArg (\m -> set_match $ MatchNotRegex $ Text.Regex.mkRegexWithOpts m False False) "REGEX")
      ("perform following action for each line not matching REGEX (case-insensitive)"),

    GetOpt.Option "h" ["header"]
      (GetOpt.ReqArg (\h -> add_action $ ActionSubst (h ++ "\\_")) "TEXT")
      ("display TEXT header before (matching) lines (same as -s 'TEXT\\_')"),
    GetOpt.Option "c" ["color"]
      (GetOpt.ReqArg (\c -> add_action $ ActionColor $ parseColor c) "COLOR")
      ("display (matching) lines in COLOR (valid colors are: normal, br,ul,bl,rev,hidden, nobr,noul..., black,red,green,yellow,blue,magenta,cyan,white, /black,/red,...)"),
    GetOpt.Option "d" ["hide"]
      (GetOpt.NoArg (add_action ActionHide))
      ("hide (matching) lines"),
    GetOpt.Option "s" ["substitute"]
      (GetOpt.ReqArg (\s -> add_action $ ActionSubst s) "TEXT")
      ("substitute (matching) lines with TEXT"),
    GetOpt.Option "e" ["execute"]
      (GetOpt.ReqArg (\e -> add_action $ ActionExecute e) "PROG")
      ("execute PROG for every (matching) line")
  ]
prog_arg a Options{ optionTails = l, optionTail = t } =
  Options{
    optionTails = t{
      tailTarg = read a,
      tailMatches = reverse (tailMatches t)
    } : l,
    optionTail = t{
      tailBegin = False,
      tailMatches = []
    },
    optionMatch = MatchAll
  }

run :: [Tail] -> IO (Control.Concurrent.MVar ExitCode)
run tails = do
  emv <- Control.Concurrent.newEmptyMVar
  let exit = Control.Concurrent.putMVar emv
  let exit0 = exit ExitSuccess
  let exit1 = exit (ExitFailure 1)
  installHandler sigINT (CatchOnce exit0) Nothing

  lockv <- Control.Concurrent.newQSem 0
  let lock = Control.Exception.bracket_ (Control.Concurrent.waitQSem lockv) (Control.Concurrent.signalQSem lockv)
  let unlock = Control.Exception.bracket_ (Control.Concurrent.signalQSem lockv) (Control.Concurrent.waitQSem lockv)

  count <- Data.IORef.newIORef (length tails)
#ifdef INOTIFY
  inotify <- 
    catchOnlyIOET GHC.IOBase.UnsupportedOperation 
      (Just =.< INotify.initINotify) 
      (return Nothing)
#endif

  let error t e = if isint e then exit0 else tailErrMsg t (show e) >> exit1
      runt = runTail TailRuntime
	{ trText = tailText
	, trUnlock = unlock
#ifdef INOTIFY
	, trINotify = inotify
#endif
	}
  mapM_ (\t -> Control.Concurrent.forkIO $ lock $ Control.Exception.handle (error t) $ do
      runt t
      i <- Data.IORef.atomicModifyIORef count (\i -> (pred i, i))
      when (i == 1) exit0)
    tails
  Control.Concurrent.signalQSem lockv
  return emv
  where
    isint (Control.Exception.IOException e) = GHC.IOBase.ioe_type e == GHC.IOBase.Interrupted
    isint _ = False

main = do
  args <- System.Environment.getArgs
  tails <- case (GetOpt.getOpt (GetOpt.ReturnInOrder prog_arg) prog_options args) of
    (s, [], []) -> case optionTails $ foldl (\s t -> t s) defaultOptions s of
      [] -> do
	putStrLn prog_usage
	exitWith ExitSuccess
      t -> return t
    (_, _, err) -> do
      mapM_ putStrLn err
      putStrLn prog_usage
      exitFailure
  e <- run tails >>= Control.Concurrent.takeMVar
  when (e == ExitSuccess) $
    errMsg "ztail: done"
  exitWith e