agitprop.hs

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