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