agitprop.hs

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