agitprop.hs

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