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