agitprop.hs

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