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
  45       printparam "hotfolder"    (inputFile   args) >>
  46       printparam "initial scan" (initialScan args)
  47     else printparam
  48            "input file"
  49            [ inputFile args
  50            , if (lineMode args)
  51                then "(line-by-line)"
  52                else ""
  53            ]
  54   printparam "remove sent file" (removeSentFile args && isDir)
  55   (conn, chan) <- connect args
  56   addChannelExceptionHandler chan (X.throwTo tid)
  57   printparam "confirm mode" $ confirm args
  58   if (confirm args)
  59     then do
  60       confirmSelect chan False
  61       addConfirmationListener chan confirmCallback
  62     else return ()
  63   let publishOneMsg = 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 "." >>= mapM_ (\fn -> handleFile publishOneMsg (suffix args) fn)
  70            else return()
  71          inotify <- initINotify
  72          wd <-
  73            addWatch
  74              inotify
  75              [CloseWrite, MoveIn]
  76              "."
  77              (handleEvent publishOneMsg (suffix args))
  78          hr $ "BEGIN watching " ++ (inputFile args)
  79          sleepingBeauty >>= printparam "exception"
  80          removeWatch wd
  81          hr $ "END watching " ++ (inputFile args)
  82        else do
  83          hr $ "BEGIN sending"
  84          messageFile <-
  85            if inputFile args == "-"
  86              then BL.getContents
  87              else BL.readFile (inputFile args)
  88          if (lineMode args)
  89            then mapM_ (publishOneMsg Nothing) (BL.lines messageFile)
  90            else publishOneMsg (Just (inputFile args)) messageFile
  91          hr "END sending")
  92     exceptionHandler
  93   -- all done. wait and close.
  94   if (confirm args)
  95     then waitForConfirms chan >>= printparam "confirmed"
  96     else return ()
  97   X.catch (closeConnection conn) exceptionHandler
  98
  99 -- | A handler for clean exit
 100 exceptionHandler :: AMQPException -> IO ()
 101 exceptionHandler (ChannelClosedException Normal txt) =
 102   printparam "exit" txt >> exitWith ExitSuccess
 103 exceptionHandler (ConnectionClosedException Normal txt) =
 104   printparam "exit" txt >> exitWith ExitSuccess
 105 exceptionHandler x = printparam "exception" x >> exitWith (ExitFailure 1)
 106
 107 -- | The handler for publisher confirms
 108 confirmCallback :: (Word64, Bool, AckType) -> IO ()
 109 confirmCallback (deliveryTag, isAll, ackType) =
 110   printparam
 111     "confirmed"
 112     [ show deliveryTag
 113     , if isAll
 114         then "all"
 115         else "this"
 116     , show ackType
 117     ]
 118
 119 -- | Hotfolder event handler
 120 handleEvent ::
 121      (Maybe String -> BL.ByteString -> IO ())
 122   -> [String]
 123   -> Event
 124   -> IO ()
 125 -- just handle closewrite and movedin events
 126 #if MIN_VERSION_hinotify(0,3,10)
 127 handleEvent func suffixes (Closed False (Just fileName) True) =
 128   handleFile func suffixes (BS.unpack fileName)
 129 handleEvent func suffixes (MovedIn False fileName _) =
 130   handleFile func suffixes (BS.unpack fileName)
 131 #else
 132 handleEvent func suffixes (Closed False (Just fileName) True) = handleFile func suffixes fileName
 133 handleEvent func suffixes (MovedIn False fileName _) = 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) = printparam "removing" fname >> removeFile fname