agitprop.hs

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