agitprop.hs

   1 {-# LANGUAGE CPP               #-}
   2 {-# LANGUAGE OverloadedStrings #-}
   3
   4 -- generic AMQP publisher
   5 import           Control.Concurrent
   6 import qualified Control.Exception             as X
   7 import qualified Data.ByteString.Lazy.Char8    as BL
   8 #if linux_HOST_OS
   9 #if MIN_VERSION_hinotify(0,3,10)
  10 import qualified Data.ByteString.Char8         as BS
  11 #endif
  12 import           Data.List                     (isSuffixOf)
  13 #endif
  14 import           Data.Maybe
  15 import qualified Data.Text                     as T
  16 import           Data.Time
  17 import           Data.Time.Clock.POSIX
  18 import           Data.Version                  (showVersion)
  19 import           Data.Word                     (Word64)
  20 import           Magic
  21 import           Network.AMQP
  22 import           Network.AMQP.Types
  23 import           Network.AMQP.Utils.Connection
  24 import           Network.AMQP.Utils.Helpers
  25 import           Network.AMQP.Utils.Options
  26 import           Paths_amqp_utils              (version)
  27 import           System.Directory
  28 import           System.Environment
  29 import           System.Exit
  30 #if linux_HOST_OS
  31 import           System.INotify
  32 #endif
  33 import qualified System.Posix.Files            as F
  34
  35 main :: IO ()
  36 main = do
  37   hr "starting"
  38   tid <- myThreadId
  39   args <- getArgs >>= parseargs 'a'
  40   printparam "client version" ["amqp-utils", showVersion version]
  41   printparam "routing key" $ rKey args
  42   printparam "exchange" $ currentExchange args
  43   isDir <-
  44     if inputFile args == "-"
  45       then return False
  46       else F.getFileStatus (inputFile args) >>= return . F.isDirectory
  47   if isDir
  48     then printparam "hotfolder" (inputFile args) >>
  49          printparam "initial scan" (initialScan args)
  50     else printparam
  51            "input file"
  52            [ inputFile args
  53            , if (lineMode args)
  54                then "(line-by-line)"
  55                else ""
  56            ]
  57   printparam "remove sent file" (removeSentFile args && isDir)
  58   (conn, chan) <- connect args
  59   addChannelExceptionHandler chan (X.throwTo tid)
  60   printparam "confirm mode" $ confirm args
  61   if (confirm args)
  62     then do
  63       confirmSelect chan False
  64       addConfirmationListener chan confirmCallback
  65     else return ()
  66   let publishOneMsg =
  67         publishOneMsg' chan args {removeSentFile = removeSentFile args && isDir}
  68   X.catch
  69     (if isDir
  70        then do
  71 #if linux_HOST_OS
  72          setCurrentDirectory (inputFile args)
  73          if (initialScan args)
  74            then getDirectoryContents "." >>=
  75                 mapM_ (\fn -> handleFile publishOneMsg (suffix args) fn)
  76            else return ()
  77          inotify <- initINotify
  78          wd <-
  79            addWatch
  80              inotify
  81              [CloseWrite, MoveIn]
  82              "."
  83              (handleEvent publishOneMsg (suffix args))
  84          hr $ "BEGIN watching " ++ (inputFile args)
  85          sleepingBeauty >>= printparam "exception"
  86          removeWatch wd
  87          hr $ "END watching " ++ (inputFile args)
  88 #else
  89          X.throw (X.ErrorCall "ERROR: watching a directory is only supported in Linux")
  90 #endif
  91        else do
  92          hr $ "BEGIN sending"
  93          messageFile <-
  94            if inputFile args == "-"
  95              then BL.getContents
  96              else BL.readFile (inputFile args)
  97          if (lineMode args)
  98            then mapM_ (publishOneMsg Nothing) (BL.lines messageFile)
  99            else publishOneMsg (Just (inputFile args)) messageFile
 100          hr "END sending")
 101     exceptionHandler
 102   -- all done. wait and close.
 103   if (confirm args)
 104     then waitForConfirms chan >>= printparam "confirmed"
 105     else return ()
 106   X.catch (closeConnection conn) exceptionHandler
 107
 108 -- | A handler for clean exit
 109 exceptionHandler :: AMQPException -> IO ()
 110 exceptionHandler (ChannelClosedException Normal txt) =
 111   printparam "exit" txt >> exitWith ExitSuccess
 112 exceptionHandler (ConnectionClosedException Normal txt) =
 113   printparam "exit" txt >> exitWith ExitSuccess
 114 exceptionHandler x = printparam "exception" x >> exitWith (ExitFailure 1)
 115
 116 -- | The handler for publisher confirms
 117 confirmCallback :: (Word64, Bool, AckType) -> IO ()
 118 confirmCallback (deliveryTag, isAll, ackType) =
 119   printparam
 120     "confirmed"
 121     [ show deliveryTag
 122     , if isAll
 123         then "all"
 124         else "this"
 125     , show ackType
 126     ]
 127
 128 #if linux_HOST_OS
 129 -- | Hotfolder event handler
 130 handleEvent ::
 131      (Maybe String -> BL.ByteString -> IO ()) -> [String] -> Event -> IO ()
 132 -- just handle closewrite and movedin events
 133 #if MIN_VERSION_hinotify(0,3,10)
 134 handleEvent func suffixes (Closed False (Just fileName) True) =
 135   handleFile func suffixes (BS.unpack fileName)
 136 handleEvent func suffixes (MovedIn False fileName _) =
 137   handleFile func suffixes (BS.unpack fileName)
 138 #else
 139 handleEvent func suffixes (Closed False (Just fileName) True) =
 140   handleFile func suffixes fileName
 141 handleEvent func suffixes (MovedIn False fileName _) =
 142   handleFile func suffixes fileName
 143 #endif
 144 handleEvent _ _ _ = return ()
 145
 146 -- | Hotfolder file handler
 147 handleFile ::
 148      (Maybe String -> BL.ByteString -> IO ()) -> [String] -> FilePath -> IO ()
 149 handleFile _ _ ('.':_) = return () -- ignore hidden files
 150 handleFile func suffixes@(_:_) fileName =
 151   if any (flip isSuffixOf fileName) suffixes
 152     then handleFile func [] fileName
 153     else return ()
 154 handleFile func [] fileName =
 155   X.catch
 156     (BL.readFile fileName >>= func (Just fileName))
 157     (\e -> printparam "exception in handleFile" (e :: X.IOException))
 158 #endif
 159
 160 -- | Publish one message with our settings
 161 publishOneMsg' :: Channel -> Args -> Maybe FilePath -> BL.ByteString -> IO ()
 162 publishOneMsg' chan a fn content = do
 163   printparam "sending" fn
 164   (mtype, mencoding) <-
 165     if (magic a) && isJust fn
 166       then do
 167         m <- magicOpen [MagicMimeType]
 168         magicLoadDefault m
 169         t <- magicFile m (fromJust fn)
 170         magicSetFlags m [MagicMimeEncoding]
 171         e <- magicFile m (fromJust fn)
 172         return (Just (T.pack t), Just (T.pack e))
 173       else return ((contenttype a), (contentencoding a))
 174   now <- getCurrentTime >>= return . floor . utcTimeToPOSIXSeconds
 175   publishMsg
 176     chan
 177     (T.pack $ currentExchange a)
 178     (T.pack $ rKey a)
 179     newMsg
 180       { msgBody = content
 181       , msgDeliveryMode = persistent a
 182       , msgTimestamp = Just now
 183       , msgID = msgid a
 184       , msgType = msgtype a
 185       , msgUserID = userid a
 186       , msgApplicationID = appid a
 187       , msgClusterID = clusterid a
 188       , msgContentType = mtype
 189       , msgContentEncoding = mencoding
 190       , msgReplyTo = replyto a
 191       , msgPriority = prio a
 192       , msgCorrelationID = corrid a
 193       , msgExpiration = msgexp a
 194       , msgHeaders = substheader (fnheader a) fn $ msgheader a
 195       } >>=
 196     printparam "sent"
 197   removeSentFileIfRequested (removeSentFile a) fn
 198   where
 199     substheader ::
 200          [String] -> Maybe String -> Maybe FieldTable -> Maybe FieldTable
 201     substheader (s:r) (Just fname) old =
 202       substheader r (Just fname) (addheader old (s ++ "=" ++ fname))
 203     substheader _ _ old = old
 204     removeSentFileIfRequested False _ = return ()
 205     removeSentFileIfRequested True Nothing = return ()
 206     removeSentFileIfRequested True (Just fname) =
 207       printparam "removing" fname >> removeFile fname