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