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
  46       printparam "hotfolder"    (inputFile   args) >>
  47       printparam "initial scan" (initialScan args)
  48     else printparam
  49            "input file"
  50            [ inputFile args
  51            , if (lineMode args)
  52                then "(line-by-line)"
  53                else ""
  54            ]
  55   printparam "remove sent file" (removeSentFile args && isDir)
  56   (conn, chan) <- connect args
  57   addChannelExceptionHandler chan (X.throwTo tid)
  58   printparam "confirm mode" $ confirm args
  59   if (confirm args)
  60     then do
  61       confirmSelect chan False
  62       addConfirmationListener chan confirmCallback
  63     else return ()
  64   let publishOneMsg = publishOneMsg' chan args {removeSentFile = removeSentFile args && isDir}
  65   X.catch
  66     (if isDir
  67        then do
  68          setCurrentDirectory (inputFile args)
  69          if (initialScan args)
  70            then getDirectoryContents "." >>= 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          _ <- forever $ threadDelay 1000000
  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 ())
 123   -> [String]
 124   -> Event
 125   -> IO ()
 126 -- just handle closewrite and movedin events
 127 #if MIN_VERSION_hinotify(0,3,10)
 128 handleEvent func suffixes (Closed False (Just fileName) True) =
 129   handleFile func suffixes (BS.unpack fileName)
 130 handleEvent func suffixes (MovedIn False fileName _) =
 131   handleFile func suffixes (BS.unpack fileName)
 132 #else
 133 handleEvent func suffixes (Closed False (Just fileName) True) = handleFile func suffixes fileName
 134 handleEvent func suffixes (MovedIn False fileName _) = handleFile func suffixes fileName
 135 #endif
 136 handleEvent _ _ _ = return ()
 137
 138 -- | Hotfolder file handler
 139 handleFile ::
 140      (Maybe String -> BL.ByteString -> IO ()) -> [String] -> FilePath -> IO ()
 141 handleFile _ _ ('.':_) = return () -- ignore hidden files
 142 handleFile func suffixes@(_:_) fileName =
 143   if any (flip isSuffixOf fileName) suffixes
 144     then handleFile func [] fileName
 145     else return ()
 146 handleFile func [] fileName =
 147   X.catch
 148     (BL.readFile fileName >>= func (Just fileName))
 149     (\e -> printparam "exception in handleFile" (e :: X.SomeException))
 150
 151 -- | Publish one message with our settings
 152 publishOneMsg' :: Channel -> Args -> Maybe FilePath -> BL.ByteString -> IO ()
 153 publishOneMsg' chan a fn content = do
 154   printparam "sending" fn
 155   (mtype, mencoding) <-
 156     if (magic a) && isJust fn
 157       then do
 158         m <- magicOpen [MagicMimeType]
 159         magicLoadDefault m
 160         t <- magicFile m (fromJust fn)
 161         magicSetFlags m [MagicMimeEncoding]
 162         e <- magicFile m (fromJust fn)
 163         return (Just (T.pack t), Just (T.pack e))
 164       else return ((contenttype a), (contentencoding a))
 165   now <- getCurrentTime >>= return . floor . utcTimeToPOSIXSeconds
 166   publishMsg
 167     chan
 168     (T.pack $ currentExchange a)
 169     (T.pack $ rKey a)
 170     newMsg
 171       { msgBody = content
 172       , msgDeliveryMode = persistent a
 173       , msgTimestamp = Just now
 174       , msgID = msgid a
 175       , msgType = msgtype a
 176       , msgUserID = userid a
 177       , msgApplicationID = appid a
 178       , msgClusterID = clusterid a
 179       , msgContentType = mtype
 180       , msgContentEncoding = mencoding
 181       , msgReplyTo = replyto a
 182       , msgPriority = prio a
 183       , msgCorrelationID = corrid a
 184       , msgExpiration = msgexp a
 185       , msgHeaders = substheader (fnheader a) fn $ msgheader a
 186       } >>=
 187     printparam "sent"
 188   removeSentFileIfRequested (removeSentFile a) fn
 189   where
 190     substheader ::
 191          [String] -> Maybe String -> Maybe FieldTable -> Maybe FieldTable
 192     substheader (s:r) (Just fname) old =
 193       substheader r (Just fname) (addheader old (s ++ "=" ++ fname))
 194     substheader _ _ old = old
 195     removeSentFileIfRequested False _           = return ()
 196     removeSentFileIfRequested True Nothing      = return ()
 197     removeSentFileIfRequested True (Just fname) = printparam "removing" fname >> removeFile fname