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