{-# LANGUAGE CPP #-} module Main where import System.Directory import System.FilePath import System.Console.GetOpt import System.Environment import System.Exit import System.IO import qualified Data.Map as M import Data.Version (showVersion) import Data.Maybe import Control.Monad import Control.Applicative import Data.Conduit.Attoparsec import Data.Conduit #if MIN_VERSION_aeson(2,2,0) import Data.Aeson (parseJSON) import Data.Aeson.Parser (json) #else import Data.Aeson (parseJSON, json) #endif import Data.Aeson.Types (parseEither) import Data.Binary.StringRef import Data.Attoparsec.ByteString.Char8 (skipSpace, option) import qualified Data.ByteString.Char8 as BS import qualified Data.ByteString.Lazy as BSL import qualified Data.Conduit.List as C import qualified Data.Conduit.Binary as C import TimeLog import Data import CommonStartup import DumpFormat import LockFile import Paths_arbtt (version) data Options = Options { optLogFile :: String , optFormat :: DumpFormat , optAppend :: Bool } defaultOptions dir = Options { optLogFile = dir "capture.log" , optAppend = False , optFormat = DFShow } versionStr = "arbtt-import " ++ showVersion version header = "Usage: arbtt-import [OPTIONS...]" options :: [OptDescr (Options -> IO Options)] options = [ Option "h?" ["help"] (NoArg $ \_ -> do putStr (usageInfo header options) exitSuccess ) "show this help" , Option "V" ["version"] (NoArg $ \_ -> do putStrLn versionStr exitSuccess ) "show the version number" , Option "f" ["logfile"] (ReqArg (\arg opt -> return opt { optLogFile = arg }) "FILE") "use this file instead of ~/.arbtt/capture.log" , Option "a" ["append"] (NoArg (\opt -> return opt { optAppend = True } )) "append to the logfile, instead of overriding it" , Option "t" ["format"] (ReqArg (\arg opt -> case readDumpFormat arg of Just DFHuman -> do hPutStrLn stderr "The “Human” format cannot imported." hPutStr stderr (usageInfo header options) exitFailure Just fm -> return $ opt { optFormat = fm} Nothing -> do hPutStrLn stderr ("Invalid format \"" ++ arg ++ "\".") hPutStr stderr (usageInfo header options) exitFailure ) "FORMAT") "output format, one of Show (default) or JSON " ] parseConduit :: DumpFormat -> Conduit BS.ByteString IO [TimeLogEntry CaptureData] parseConduit DFHuman = error "Cannot read back human format" parseConduit DFShow = C.lines =$= C.map ( (:[]) . read . BS.unpack) parseConduit DFJSON = conduitParser (ignoreWhiteSpace json) =$= C.map snd =$= C.catMaybes =$= tlConduit where tlConduit = C.mapM $ \ v -> case parseEither oneOrMany v of Left e -> do hPutStrLn stderr ("Cannot parse log entry: " ++ e) exitFailure Right x -> pure x oneOrMany v = parseJSON v <|> ((:[]) <$> parseJSON v) ignoreWhiteSpace p = skipSpace *> option Nothing (Just <$> p) binaryConduit :: ListOfStringable a => Conduit [(TimeLogEntry a, Maybe a)] IO (Flush BS.ByteString) binaryConduit = C.map (map go) =$= C.concatMap chunk where go (x,prev) = BSL.toStrict $ ls_encode strs x where strs = maybe [] listOfStrings prev chunk xs = map Chunk xs ++ [Flush] main = do commonStartup args <- getArgs actions <- case getOpt Permute options args of (o,[],[]) -> return o (_,_,errs) -> do hPutStr stderr (concat errs ++ usageInfo header options) exitFailure dir <- getAppUserDataDirectory "arbtt" flags <- foldl (>>=) (return (defaultOptions dir)) actions lockFile (optLogFile flags) ex <- doesFileExist (optLogFile flags) when (ex && not (optAppend flags)) $ do putStrLn $ "File at " ++ (optLogFile flags) ++ " does already exist. Please delete this" putStrLn $ "file before running arbtt-import." exitFailure createTimeLog False (optLogFile flags) h <- openBinaryFile (optLogFile flags) AppendMode runConduit $ C.sourceHandle stdin =$= parseConduit (optFormat flags) =$= stutterList =$= binaryConduit =$= C.sinkHandleFlush h hClose h stutter :: Monad m => Conduit (TimeLogEntry a) m (TimeLogEntry a, Maybe a) stutter = loop Nothing where loop prev = do mx <- await case mx of Nothing -> return () Just x -> do yield (x, prev) loop (Just (tlData x)) stutterList :: Monad m => Conduit [TimeLogEntry a] m [(TimeLogEntry a, Maybe a)] stutterList = loop Nothing where loop prev = do mx <- await case mx of Nothing -> return () Just xs -> do let xs' = zip xs (prev : map (Just . tlData) xs) yield xs' loop (Just (tlData (last xs)))