agitprop.hs

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