ztail-1.2: Main.hs

{-# LANGUAGE CPP, TupleSections, OverloadedStrings #-}
import Control.Concurrent.MVar (MVar, newEmptyMVar, putMVar, takeMVar)
import Control.Exception (AsyncException(UserInterrupt), fromException, handle)
import Control.Monad (void, when, join)
import Data.Bits ((.|.))
import qualified Data.ByteString.Char8 as BS
import Data.IORef (newIORef, readIORef, atomicModifyIORef')
import Data.List (foldl')
import Data.Monoid ((<>))
import GHC.IO.Exception (IOErrorType(UnsupportedOperation))
import qualified System.Console.GetOpt as Opt
import System.Environment (getArgs)
import System.Exit (ExitCode(..), exitSuccess, exitFailure, exitWith)
#ifdef INOTIFY
import qualified System.INotify as INotify
#endif
import System.IO.Error (ioeGetErrorType)
import System.Posix.Signals (installHandler, sigINT, Handler(..))
import Text.Regex.Posix (makeRegexOpts, compExtended, compIgnoreCase, compNoSub, defaultExecOpt)

import Util
import Display
import TailTypes
import TailHandle

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

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

defaultOptions :: Options
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, prog_usage :: String
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 those marked '*' which apply to all following FILEs.  Match options\n\
(-amn) apply to all following actions (-hcdse).  Actions involving TEXT can\n\
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 = Opt.usageInfo prog_header prog_options

prog_options :: [Opt.OptDescr (Options -> Options)]
prog_options = 
  [ Opt.Option "i" ["interval"]
      (Opt.ReqArg (\i -> set_opt $ \p -> p
        { tailPollInterval = read i
#ifdef INOTIFY
        , tailPollINotify = False
#endif
        }) "INT")
      ("*poll for data every INT seconds [" ++ show (tailPollInterval defaultTail) ++ "]")
  , Opt.Option "r" ["reopen"]
      (Opt.OptArg (\i -> set_opt $ \p -> p
        { tailReopenInterval = maybe (tailPollInterval p) read i
#ifdef INOTIFY
        , tailReopenINotify = False
#endif
        }) "INT")
      ("*check file name (like tail -F) every INT seconds or every poll [" ++ show (tailReopenInterval defaultTail) ++ "]")
#ifdef INOTIFY
  , Opt.Option "I" ["inotify"]
      (Opt.OptArg (\i -> set_opt $ \p -> p
        { tailPollINotify = True
        , tailPollInterval = maybe 0 read i }) "INT")
      ("*use inotify to poll for new data (and also poll every INT)")
  , Opt.Option "R" ["ireopen"]
      (Opt.NoArg (set_opt $ \p -> p
        { tailReopenINotify = True }))
      ("*use inotify to monitor file renames (only for preexisting, leaf files)")
#endif
  , Opt.Option "b" ["begin"]
      (Opt.NoArg (set_opt $ \p -> p
        { tailBegin = True }))
      (" start reading at the beginning of the file (rather than only new lines at the end)")
  , Opt.Option "l" ["dirlist"]
      (Opt.NoArg (set_opt $ \p -> p
        { tailDirList = True }))
      (" watch the contents of a directory, reporting when files are added or removed")
  , Opt.Option "D" ["dirtail"]
      (Opt.NoArg (set_opt $ \p -> p
        { tailDirTail = True }))
      (" tail all the files in a directory")
  , Opt.Option "A" ["recursive"]
      (Opt.NoArg (set_opt $ \p -> p
        { tailDirRecursive = True }))
      (" apply the above directory modifiers recursively")

  , Opt.Option "t" ["timefmt"]
      (Opt.ReqArg (\t -> set_opt $ \p -> p
        { tailTimeFmt = t }) "FMT")
      ("*set time format for \\@ substitution (in strftime(3)) [" ++ tailTimeFmt defaultTail ++ "]")
  , Opt.Option "T" ["timestamp"]
      (Opt.OptArg (maybe id $ \t -> add_action (ActionSubst "\\@ \\_") . set_opt (\p -> p
        { tailTimeFmt = t })) "FMT")
      (" timestamp with FMT; equivalent to: [-t FMT] -h '\\@ '")

  , Opt.Option "a" ["all"]
      (Opt.NoArg (set_match MatchAll))
      (" perform following action for every line from this FILE (default)")
  , Opt.Option "m" ["match"]
      (Opt.ReqArg (set_match . MatchRegex . makeRegexOpts compExtended defaultExecOpt) "REGEX")
      (" perform following action for each line matching REGEX")
  , Opt.Option "M" ["imatch"]
      (Opt.ReqArg (set_match . MatchRegex . makeRegexOpts (compExtended .|. compIgnoreCase) defaultExecOpt) "REGEX")
      (" perform following action for each line matching REGEX (case-insensitive)")
  , Opt.Option "n" ["no-match"]
      (Opt.ReqArg (set_match . MatchNotRegex . makeRegexOpts (compExtended .|. compNoSub) defaultExecOpt) "REGEX")
      (" perform following action for each line not matching REGEX")
  , Opt.Option "N" ["no-imatch"]
      (Opt.ReqArg (set_match . MatchNotRegex . makeRegexOpts (compExtended .|. compNoSub .|. compIgnoreCase) defaultExecOpt) "REGEX")
      (" perform following action for each line not matching REGEX (case-insensitive)")

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

run :: [Tail] -> IO (MVar ExitCode)
run tails = do
  emv <- newEmptyMVar
  count <- newIORef (length tails)
  errors <- newIORef False
#ifdef INOTIFY
  inotify <- 
    catchWhen ((UnsupportedOperation ==) . ioeGetErrorType)
      (Just <$> INotify.initINotify) 
      (return Nothing)
#endif
  out <- runOutput
  let done = do
        e <- readIORef errors
        putMVar emv $ if e
          then ExitFailure 1
          else ExitSuccess
      err t e = case fromException e of
	Just UserInterrupt -> done
	_ -> tailErrMsg tr t (BS.pack $ show e) >> atomicModifyIORef' errors (const (True, ()))
      tr = TailRuntime
	{ trOutput = out
        , trAddTail = (atomicModifyIORef' count ((, ()) . succ) >>) . runt
#ifdef INOTIFY
	, trINotify = inotify
#endif
	}
      runt t = void $ forkIOUnmasked $ do
        handle (err t) $ runTail tr t
        i <- atomicModifyIORef' count (join (,) . pred)
        when (i == 0) $ done

  _ <- installHandler sigINT (CatchOnce done) Nothing
  mapM_ runt tails
  return emv

main :: IO ()
main = do
  args <- getArgs
  tails <- case Opt.getOpt (Opt.ReturnInOrder prog_arg) prog_options args of
    (s, [], []) -> case optionTails $ foldl' (flip ($)) defaultOptions s of
      [] -> do
	putStrLn prog_usage
	exitSuccess
      t -> return $ reverse t
    (_, _, err) -> do
      mapM_ putStrLn err
      putStrLn prog_usage
      exitFailure
  e <- run tails >>= takeMVar
  when (e == ExitSuccess) $
    rawErrMsg "ztail: done"
  exitWith e