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