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