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 >>= printparam "exception"
  80          forM_ wds (\(wd,folder) -> do
  81            removeWatch wd
  82            printparam "END watching" folder
  83            )
  84 #else
  85          X.throw (X.ErrorCall "ERROR: watching a directory is only supported in Linux")
  86 #endif
  87        else do
  88          hr $ "BEGIN sending"
  89          messageFile <-
  90            if inputFile' == "-"
  91              then BL.getContents
  92              else readFileRawLazy inputFile'
  93          if (lineMode args)
  94            then mapM_ (publishOneMsg (currentExchange args) (rKey args) Nothing) (BL.lines messageFile)
  95            else publishOneMsg (currentExchange args) (rKey args) (Just (inputFile')) messageFile
  96          hr "END sending")
  97     exceptionHandler
  98   -- all done. wait and close.
  99   if (confirm args)
 100     then waitForConfirms chan >>= printparam "confirmed"
 101     else return ()
 102   X.catch (closeConnection conn) exceptionHandler
 103
 104 #if linux_HOST_OS
 105 -- | watch a hotfolder
 106 watchHotfolder ::
 107      Args
 108   -> (String -> String -> Maybe RawFilePath -> BL.ByteString -> IO ())
 109   -> (RawFilePath, String, String)
 110   -> IO (WatchDescriptor, RawFilePath)
 111 watchHotfolder args publishOneMsg (folder, exchange, rkey) = do
 112   printparam "hotfolder" folder
 113   inotify <- initINotify
 114   wd <-
 115    addWatch
 116      inotify
 117      [CloseWrite, MoveIn]
 118      folder
 119      (handleEvent (publishOneMsg exchange rkey) (suffix args) folder)
 120   hr "BEGIN watching"
 121   if (initialScan args)
 122    then RD.listDirectory folder >>=
 123         mapM_ (\fn -> handleFile (publishOneMsg exchange rkey) (suffix args) (folder </> fn))
 124    else return ()
 125   return (wd,folder)
 126 #endif
 127
 128 -- | A handler for clean exit
 129 exceptionHandler :: AMQPException -> IO ()
 130 exceptionHandler (ChannelClosedException Normal txt) =
 131   printparam "exit" txt >> exitWith ExitSuccess
 132 exceptionHandler (ConnectionClosedException Normal txt) =
 133   printparam "exit" txt >> exitWith ExitSuccess
 134 exceptionHandler x = printparam "exception" x >> exitWith (ExitFailure 1)
 135
 136 -- | The handler for publisher confirms
 137 confirmCallback :: (Word64, Bool, AckType) -> IO ()
 138 confirmCallback (deliveryTag, isAll, ackType) =
 139   printparam
 140     "confirmed"
 141     [ show deliveryTag
 142     , if isAll
 143         then "all"
 144         else "this"
 145     , show ackType
 146     ]
 147 #if linux_HOST_OS
 148 -- | Hotfolder event handler
 149 handleEvent ::
 150      (Maybe RawFilePath -> BL.ByteString -> IO ()) -> [BS.ByteString] -> RawFilePath -> Event -> IO ()
 151 -- just handle closewrite and movedin events
 152 handleEvent func suffixes folder (Closed False (Just fileName) True) =
 153   handleFile func suffixes (folder </> fileName)
 154 handleEvent func suffixes folder (MovedIn False fileName _) =
 155   handleFile func suffixes (folder </> fileName)
 156 handleEvent _ _ _ _ = return ()
 157
 158 -- | Hotfolder file handler
 159 handleFile ::
 160      (Maybe RawFilePath -> BL.ByteString -> IO ()) -> [BS.ByteString] -> RawFilePath -> IO ()
 161 handleFile func suffixes@(_:_) fileName =
 162   if (any (flip BS.isSuffixOf fileName) suffixes) && not ("." `BS.isPrefixOf` fileName)
 163     then handleFile func [] fileName
 164     else return ()
 165 handleFile func [] fileName =
 166   X.catch
 167     (readFileRawLazy fileName >>= func (Just fileName))
 168     (\e -> printparam "exception while processing" fileName >> printparam "exception in handleFile" (e :: X.IOException))
 169 #endif
 170
 171 -- | Publish one message with our settings
 172 publishOneMsg' ::
 173      Channel
 174   -> Args
 175   -> String
 176   -> String
 177   -> Maybe RawFilePath
 178   -> BL.ByteString
 179   -> IO ()
 180 publishOneMsg' chan a exchange rkey fn content = do
 181   printparam "sending" fn
 182   (mtype, mencoding) <-
 183     if (magic a)
 184       then do
 185         let firstchunk = if BL.null content then BS.empty else head $ BL.toChunks content
 186         m <- magicOpen [MagicMimeType]
 187         magicLoadDefault m
 188         t <- BS.useAsCStringLen firstchunk (magicCString m)
 189         magicSetFlags m [MagicMimeEncoding]
 190         e <- BS.useAsCStringLen firstchunk (magicCString m)
 191         return (Just (T.pack t), Just (T.pack e))
 192       else return ((contenttype a), (contentencoding a))
 193   now <- getCurrentTime >>= return . floor . utcTimeToPOSIXSeconds
 194   publishMsg
 195     chan
 196     (T.pack $ exchange)
 197     (T.pack $ rkey)
 198     newMsg
 199       { msgBody = content
 200       , msgDeliveryMode = persistent a
 201       , msgTimestamp = Just now
 202       , msgID = msgid a
 203       , msgType = msgtype a
 204       , msgUserID = userid a
 205       , msgApplicationID = appid a
 206       , msgClusterID = clusterid a
 207       , msgContentType = mtype
 208       , msgContentEncoding = mencoding
 209       , msgReplyTo = replyto a
 210       , msgPriority = prio a
 211       , msgCorrelationID = corrid a
 212       , msgExpiration = msgexp a
 213       , msgHeaders = substheader (fnheader a) (fmap takeFileName fn) $ msgheader a
 214       } >>=
 215     printparam "sent"
 216   removeSentFileIfRequested (removeSentFile a) (moveSentFileTo a) fn
 217   where
 218     substheader ::
 219          [String] -> Maybe BS.ByteString -> Maybe FieldTable -> Maybe FieldTable
 220     substheader (s:r) (Just fname) old =
 221       substheader r (Just fname) (addheader' old s fname)
 222     substheader _ _ old = old
 223     removeSentFileIfRequested False _ _ = return ()
 224     removeSentFileIfRequested True _ Nothing = return ()
 225     removeSentFileIfRequested True Nothing (Just fname) =
 226       printparam "removing" fname >> RD.removeFile fname
 227     removeSentFileIfRequested True (Just path) (Just fname) =
 228       printparam "moving" [fname,"to",path] >>
 229       F.rename fname (replaceDirectory fname ((takeDirectory fname) </> path))
 230     addheader' :: Maybe FieldTable -> String -> BS.ByteString -> Maybe FieldTable
 231     addheader' Nothing k v =
 232       Just $ FieldTable $ M.singleton (T.pack k) (FVString v)
 233     addheader' (Just (FieldTable oldheader)) k v =
 234       Just $ FieldTable $ M.insert (T.pack k) (FVString v) oldheader