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                    (forM_)
   8 import qualified Data.ByteString.Lazy.Char8       as BL
   9 import qualified RawFilePath.Directory            as RD
  10 import qualified Data.ByteString.Char8            as BS
  11 import qualified Data.Map                         as M
  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.FilePath.Posix.ByteString
  29 #if linux_HOST_OS
  30 import           System.INotify
  31 #endif
  32 import qualified System.Posix.Files.ByteString    as F
  33
  34 main :: IO ()
  35 main = do
  36   hr "starting"
  37   tid <- myThreadId
  38   args <- getArgs >>= parseargs 'a'
  39   printparam "client version" ["amqp-utils", showVersion version]
  40   printparam "routing key" $ rKey args
  41   printparam "exchange" $ currentExchange args
  42   (conn, chan) <- connect args
  43   addChannelExceptionHandler chan (X.throwTo tid)
  44   printparam "confirm mode" $ confirm args
  45   if (confirm args)
  46     then do
  47       confirmSelect chan False
  48       addConfirmationListener chan confirmCallback
  49     else return ()
  50   let inputFile' = firstInputFile (inputFiles args)
  51   isDir <-
  52     if inputFile' == "-"
  53       then return False
  54       else F.getFileStatus inputFile' >>= return . F.isDirectory
  55   let publishOneMsg =
  56         publishOneMsg' chan args {removeSentFile = removeSentFile args && isDir}
  57   if isDir
  58     then printparam "initial scan" (initialScan args) >>
  59          if isNothing (moveSentFileTo args)
  60            then printparam "remove sent file" (removeSentFile args)
  61            else printparam "move sent file to" (moveSentFileTo args)
  62     else printparam
  63            "input file"
  64            [ inputFile'
  65            , if (lineMode args)
  66                then "(line-by-line)"
  67                else ""
  68            ]
  69   X.catch
  70     (if isDir
  71        then do
  72 #if linux_HOST_OS
  73          wds <- mapM (watchHotfolder args publishOneMsg) (inputFiles args)
  74          sleepingBeauty >>= printparam "exception"
  75          forM_ wds (\(wd,folder) -> do
  76            removeWatch wd
  77            hr $ "END watching " ++ folder
  78            )
  79 #else
  80          X.throw (X.ErrorCall "ERROR: watching a directory is only supported in Linux")
  81 #endif
  82        else do
  83          hr $ "BEGIN sending"
  84          messageFile <-
  85            if inputFile' == "-"
  86              then BL.getContents
  87              else readFileRawLazy inputFile'
  88          if (lineMode args)
  89            then mapM_ (publishOneMsg (currentExchange args) (rKey args) Nothing Nothing) (BL.lines messageFile)
  90            else publishOneMsg (currentExchange args) (rKey args) Nothing (Just (inputFile')) 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 #if linux_HOST_OS
 100 -- | watch a hotfolder
 101 watchHotfolder ::
 102      Args
 103   -> (String -> String -> Maybe FilePath -> Maybe RawFilePath -> BL.ByteString -> IO ())
 104   -> (FilePath, String, String)
 105   -> IO (WatchDescriptor,String)
 106 watchHotfolder args publishOneMsg (folder, exchange, rkey) = do
 107   printparam "hotfolder" folder
 108   setCurrentDirectory folder
 109   if (initialScan args)
 110    then RD.listDirectory "." >>=
 111         mapM_ (\fn -> handleFile (publishOneMsg exchange rkey (Just folder)) (suffix args) (Just folder) fn)
 112    else return ()
 113   inotify <- initINotify
 114   wd <-
 115    addWatch
 116      inotify
 117      [CloseWrite, MoveIn]
 118      (BS.pack folder)
 119      (handleEvent (publishOneMsg exchange rkey (Just folder)) (suffix args) (Just folder))
 120   hr $ "BEGIN watching " ++ folder
 121   return (wd,folder)
 122 #endif
 123
 124 -- | A handler for clean exit
 125 exceptionHandler :: AMQPException -> IO ()
 126 exceptionHandler (ChannelClosedException Normal txt) =
 127   printparam "exit" txt >> exitWith ExitSuccess
 128 exceptionHandler (ConnectionClosedException Normal txt) =
 129   printparam "exit" txt >> exitWith ExitSuccess
 130 exceptionHandler x = printparam "exception" x >> exitWith (ExitFailure 1)
 131
 132 -- | The handler for publisher confirms
 133 confirmCallback :: (Word64, Bool, AckType) -> IO ()
 134 confirmCallback (deliveryTag, isAll, ackType) =
 135   printparam
 136     "confirmed"
 137     [ show deliveryTag
 138     , if isAll
 139         then "all"
 140         else "this"
 141     , show ackType
 142     ]
 143 #if linux_HOST_OS
 144 -- | Hotfolder event handler
 145 handleEvent ::
 146      (Maybe RawFilePath -> BL.ByteString -> IO ()) -> [BS.ByteString] -> Maybe FilePath -> Event -> IO ()
 147 -- just handle closewrite and movedin events
 148 handleEvent func suffixes folder (Closed False (Just fileName) True) =
 149   handleFile func suffixes folder fileName
 150 handleEvent func suffixes folder (MovedIn False fileName _) =
 151   handleFile func suffixes folder fileName
 152 handleEvent _ _ _ _ = return ()
 153
 154 -- | Hotfolder file handler
 155 handleFile ::
 156      (Maybe RawFilePath -> BL.ByteString -> IO ()) -> [BS.ByteString] -> Maybe FilePath -> RawFilePath -> IO ()
 157 handleFile func suffixes@(_:_) folder fileName =
 158   if (any (flip BS.isSuffixOf fileName) suffixes) && not ("." `BS.isPrefixOf` fileName)
 159     then handleFile func [] folder fileName
 160     else return ()
 161 handleFile func [] folder fileName =
 162   X.catch
 163     (mapM_ setCurrentDirectory folder >> readFileRawLazy fileName >>= func (Just fileName))
 164     (\e -> printparam "exception in handleFile" (e :: X.IOException))
 165 #endif
 166
 167 -- | Publish one message with our settings
 168 publishOneMsg' ::
 169      Channel
 170   -> Args
 171   -> String
 172   -> String
 173   -> Maybe FilePath
 174   -> Maybe RawFilePath
 175   -> BL.ByteString
 176   -> IO ()
 177 publishOneMsg' chan a exchange rkey folder fn content = do
 178   printparam "sending" [fmap BS.pack folder, fn]
 179   (mtype, mencoding) <-
 180     if (magic a)
 181       then do
 182         let firstchunk = if BL.null content then BS.empty else head $ BL.toChunks content
 183         m <- magicOpen [MagicMimeType]
 184         magicLoadDefault m
 185         t <- BS.useAsCStringLen firstchunk (magicCString m)
 186         magicSetFlags m [MagicMimeEncoding]
 187         e <- BS.useAsCStringLen firstchunk (magicCString m)
 188         return (Just (T.pack t), Just (T.pack e))
 189       else return ((contenttype a), (contentencoding a))
 190   now <- getCurrentTime >>= return . floor . utcTimeToPOSIXSeconds
 191   publishMsg
 192     chan
 193     (T.pack $ exchange)
 194     (T.pack $ rkey)
 195     newMsg
 196       { msgBody = content
 197       , msgDeliveryMode = persistent a
 198       , msgTimestamp = Just now
 199       , msgID = msgid a
 200       , msgType = msgtype a
 201       , msgUserID = userid a
 202       , msgApplicationID = appid a
 203       , msgClusterID = clusterid a
 204       , msgContentType = mtype
 205       , msgContentEncoding = mencoding
 206       , msgReplyTo = replyto a
 207       , msgPriority = prio a
 208       , msgCorrelationID = corrid a
 209       , msgExpiration = msgexp a
 210       , msgHeaders = substheader (fnheader a) fn $ msgheader a
 211       } >>=
 212     printparam "sent"
 213   removeSentFileIfRequested (removeSentFile a) (fmap BS.pack (moveSentFileTo a)) fn
 214   where
 215     substheader ::
 216          [String] -> Maybe BS.ByteString -> Maybe FieldTable -> Maybe FieldTable
 217     substheader (s:r) (Just fname) old =
 218       substheader r (Just fname) (addheader' old s fname)
 219     substheader _ _ old = old
 220     removeSentFileIfRequested False _ _ = return ()
 221     removeSentFileIfRequested True _ Nothing = return ()
 222     removeSentFileIfRequested True Nothing (Just fname) =
 223       printparam "removing" fname >> RD.removeFile fname
 224     removeSentFileIfRequested True (Just path) (Just fname) =
 225       printparam "moving" [fname,"to",path] >>
 226       F.rename fname (replaceDirectory fname path)
 227     addheader' :: Maybe FieldTable -> String -> BS.ByteString -> Maybe FieldTable
 228     addheader' Nothing k v =
 229       Just $ FieldTable $ M.singleton (T.pack k) (FVString v)
 230     addheader' (Just (FieldTable oldheader)) k v =
 231       Just $ FieldTable $ M.insert (T.pack k) (FVString v) oldheader