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