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 $ "BEGIN watching " ++ (inputFile args)
  57          _ <- forever $ threadDelay 1000000
  58          removeWatch wd
  59          hr $ "END watching " ++ (inputFile args)
  60        else do
  61          hr $ "BEGIN sending"
  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          hr "END sending")
  67     (\exception -> printparam' "exception" $ show (exception :: X.SomeException))
  68   -- all done. wait and close.
  69   if (confirm args)
  70     then waitForConfirms chan >>= (printparam' "confirmed") . show
  71     else return ()
  72   closeConnection conn
  73   hr "connection closed"
  74
  75 -- | The handler for publisher confirms
  76 confirmCallback :: (Word64, Bool, AckType) -> IO ()
  77 confirmCallback (deliveryTag, isAll, ackType) =
  78   printparam'
  79     "confirmed"
  80     ((show deliveryTag) ++
  81      (if isAll
  82         then " all "
  83         else " this ") ++
  84      (show ackType))
  85
  86 -- | Hotfolder event handler
  87 handleEvent ::
  88      (Maybe String -> BL.ByteString -> IO ())
  89   -> [String]
  90   -> String
  91   -> Event
  92   -> IO ()
  93 -- just handle closewrite and movedin events
  94 handleEvent f s p (Closed False (Just x) True) = handleFile f s (p ++ "/" ++ x)
  95 handleEvent f s p (MovedIn False x _) = handleFile f s (p ++ "/" ++ x)
  96 handleEvent _ _ _ _ = return ()
  97
  98 -- | Hotfolder file handler
  99 handleFile ::
 100      (Maybe String -> BL.ByteString -> IO ()) -> [String] -> FilePath -> IO ()
 101 handleFile _ _ ('.':_) = return () -- ignore hidden files
 102 handleFile f s@(_:_) x =
 103   if any (flip isSuffixOf x) s
 104     then handleFile f [] x
 105     else return ()
 106 handleFile f [] x =
 107   X.catch
 108     (BL.readFile x >>= f (Just x))
 109     (\exception ->
 110        printparam' "exception in handleFile" $
 111        show (exception :: X.SomeException))
 112
 113 -- | Publish one message with our settings
 114 publishOneMsg' :: Channel -> Args -> Maybe String -> BL.ByteString -> IO ()
 115 publishOneMsg' c a fn f = do
 116   printparam "sending" fn
 117   (mtype, mencoding) <-
 118     if (magic a) && isJust fn
 119       then do
 120         m <- magicOpen [MagicMimeType]
 121         magicLoadDefault m
 122         t <- magicFile m (fromJust fn)
 123         magicSetFlags m [MagicMimeEncoding]
 124         e <- magicFile m (fromJust fn)
 125         return (Just (T.pack t), Just (T.pack e))
 126       else return ((contenttype a), (contentencoding a))
 127   now <- getCurrentTime >>= return . floor . utcTimeToPOSIXSeconds
 128   r <-
 129     publishMsg
 130       c
 131       (T.pack $ currentExchange a)
 132       (T.pack $ rKey a)
 133       newMsg
 134         { msgBody = f
 135         , msgDeliveryMode = persistent a
 136         , msgTimestamp = Just now
 137         , msgID = msgid a
 138         , msgType = msgtype a
 139         , msgUserID = userid a
 140         , msgApplicationID = appid a
 141         , msgClusterID = clusterid a
 142         , msgContentType = mtype
 143         , msgContentEncoding = mencoding
 144         , msgReplyTo = replyto a
 145         , msgPriority = prio a
 146         , msgCorrelationID = corrid a
 147         , msgExpiration = msgexp a
 148         , msgHeaders = substheader (fnheader a) fn $ msgheader a
 149         }
 150   printparam "sent" $ fmap show r
 151   where
 152     substheader ::
 153          [String] -> Maybe String -> Maybe FieldTable -> Maybe FieldTable
 154     substheader (s:r) (Just fname) old =
 155       substheader r (Just fname) (addheader old (s ++ "=" ++ fname))
 156     substheader _ _ old = old