agitprop.hs

   1 -- generic AMQP publisher
   2
   3 import Control.Concurrent (threadDelay)
   4 import qualified Control.Exception as X
   5 import Control.Monad (forever)
   6 import qualified Data.ByteString.Lazy.Char8 as BL
   7 import Data.List (isSuffixOf)
   8 import Data.Maybe
   9 import qualified Data.Text as T
  10 import Data.Time
  11 import Data.Time.Clock.POSIX
  12 import Data.Version (showVersion)
  13 import Data.Word (Word64)
  14 import Magic
  15 import Network.AMQP
  16 import Network.AMQP.Types
  17 import Network.AMQP.Utils.Connection
  18 import Network.AMQP.Utils.Helpers
  19 import Network.AMQP.Utils.Options
  20 import Paths_amqp_utils (version)
  21 import System.Environment
  22 import System.INotify
  23 import qualified System.Posix.Files as F
  24
  25 main :: IO ()
  26 main = do
  27   hr "starting"
  28   args <- getArgs >>= parseargs "agitprop"
  29   printparam' "client version" $ "amqp-utils " ++ (showVersion version)
  30   printparam' "routing key" $ rKey args
  31   isDir <- F.getFileStatus (inputFile args) >>= return . F.isDirectory
  32   if isDir
  33     then printparam' "hotfolder" $ inputFile args
  34     else printparam' "input file" $
  35          (inputFile args) ++
  36          if (lineMode args)
  37            then " (line-by-line)"
  38            else ""
  39   (conn, chan) <- connect args
  40   printparam' "confirm mode" $ show $ confirm args
  41   if (confirm args)
  42     then do
  43       confirmSelect chan False
  44       addConfirmationListener chan confirmCallback
  45     else return ()
  46   let publishOneMsg = publishOneMsg' chan args
  47   X.catch
  48     (if isDir
  49        then do
  50          inotify <- initINotify
  51          wd <-
  52            addWatch
  53              inotify
  54              [CloseWrite, MoveIn]
  55              (inputFile args)
  56              (handleEvent publishOneMsg (suffix args))
  57          hr $ "watching " ++ (inputFile args)
  58          _ <- forever $ threadDelay 1000000
  59          removeWatch wd
  60        else do
  61          hr $ "sending " ++ (inputFile args)
  62          messageFile <- BL.readFile (inputFile args)
  63          if (lineMode args)
  64            then mapM_ (publishOneMsg Nothing) (BL.lines messageFile)
  65            else publishOneMsg (Just (inputFile args)) messageFile)
  66     (\exception -> printparam' "exception" $ show (exception :: X.SomeException))
  67   -- all done. wait and close.
  68   if (confirm args)
  69     then waitForConfirms chan >>= return . show >> return ()
  70     else return ()
  71   closeConnection conn
  72
  73 -- | The handler for publisher confirms
  74 confirmCallback :: (Word64, Bool, AckType) -> IO ()
  75 confirmCallback (deliveryTag, isAll, ackType) =
  76   printparam'
  77     "confirmed"
  78     ((show deliveryTag) ++
  79      (if isAll
  80         then " all "
  81         else " this ") ++
  82      (show ackType))
  83
  84 -- | Hotfolder event handler
  85 handleEvent ::
  86      (Maybe String -> BL.ByteString -> IO ()) -> [String] -> Event -> IO ()
  87 -- just handle closewrite and movedin events
  88 handleEvent f s (Closed False (Just x) True) = handleFile f s x
  89 handleEvent f s (MovedIn False x _) = handleFile f s x
  90 handleEvent _ _ _ = return ()
  91
  92 -- | Hotfolder file handler
  93 handleFile ::
  94      (Maybe String -> BL.ByteString -> IO ()) -> [String] -> FilePath -> IO ()
  95 handleFile _ _ ('.':_) = return () -- ignore hidden files
  96 handleFile f s@(_:_) x =
  97   if any (flip isSuffixOf x) s
  98     then handleFile f [] x
  99     else return ()
 100 handleFile f [] x =
 101   X.catch
 102     (hr ("sending " ++ x) >> BL.readFile x >>= f (Just x))
 103     (\exception ->
 104        printparam' "exception in handleFile" $
 105        show (exception :: X.SomeException))
 106
 107 -- | Publish one message with our settings
 108 publishOneMsg' :: Channel -> Args -> Maybe String -> BL.ByteString -> IO ()
 109 publishOneMsg' c a fn f = do
 110   (mtype, mencoding) <-
 111     if (magic a) && isJust fn
 112       then do
 113         m <- magicOpen [MagicMimeType]
 114         magicLoadDefault m
 115         t <- magicFile m (fromJust fn)
 116         magicSetFlags m [MagicMimeEncoding]
 117         e <- magicFile m (fromJust fn)
 118         return (Just (T.pack t), Just (T.pack e))
 119       else return ((contenttype a), (contentencoding a))
 120   now <- getCurrentTime >>= return . floor . utcTimeToPOSIXSeconds
 121   r <-
 122     publishMsg
 123       c
 124       (T.pack $ currentExchange a)
 125       (T.pack $ rKey a)
 126       newMsg
 127         { msgBody = f
 128         , msgDeliveryMode = persistent a
 129         , msgTimestamp = Just now
 130         , msgID = msgid a
 131         , msgType = msgtype a
 132         , msgUserID = userid a
 133         , msgApplicationID = appid a
 134         , msgClusterID = clusterid a
 135         , msgContentType = mtype
 136         , msgContentEncoding = mencoding
 137         , msgReplyTo = replyto a
 138         , msgPriority = prio a
 139         , msgCorrelationID = corrid a
 140         , msgExpiration = msgexp a
 141         , msgHeaders = substheader (fnheader a) fn $ msgheader a
 142         }
 143   printparam "sent" $ fmap show r
 144   where
 145     substheader ::
 146          [String] -> Maybe String -> Maybe FieldTable -> Maybe FieldTable
 147     substheader (s:r) (Just fname) old =
 148       substheader r (Just fname) (addheader old (s ++ "=" ++ fname))
 149     substheader _ _ old = old