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.Char8            as BS
  13 import qualified Data.ByteString.Lazy.Char8       as BL
  14 import qualified Data.Map                         as M
  15 import           Data.Maybe
  16 import qualified Data.Text                        as T
  17 import           Data.Time
  18 import           Data.Time.Clock.POSIX
  19 import           Data.Version                     (showVersion)
  20 import           Data.Word                        (Word64)
  21 import           Magic
  22 import           Network.AMQP
  23 import           Network.AMQP.Types
  24 import           Network.AMQP.Utils.Connection
  25 import           Network.AMQP.Utils.Helpers
  26 import           Network.AMQP.Utils.Options
  27 import           Paths_amqp_utils                 (version)
  28 import qualified RawFilePath.Directory            as RD
  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          if (confirm args)
  99            then waitForConfirms chan >>= printparam "confirmed"
 100            else return ()
 101          X.catch (closeConnection conn) exceptionHandler
 102          )
 103     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
 149 #if linux_HOST_OS
 150 -- | Hotfolder event handler
 151 handleEvent ::
 152      (Maybe RawFilePath -> BL.ByteString -> IO ()) -> [BS.ByteString] -> RawFilePath -> Event -> IO ()
 153 -- just handle closewrite and movedin events
 154 handleEvent func suffixes folder (Closed False (Just fileName) True) =
 155   handleFile func suffixes (folder </> fileName)
 156 handleEvent func suffixes folder (MovedIn False fileName _) =
 157   handleFile func suffixes (folder </> fileName)
 158 handleEvent _ _ _ _ = return ()
 159
 160 -- | Hotfolder file handler
 161 handleFile ::
 162      (Maybe RawFilePath -> BL.ByteString -> IO ()) -> [BS.ByteString] -> RawFilePath -> IO ()
 163 handleFile func suffixes@(_:_) fileName =
 164   if (any (flip BS.isSuffixOf fileName) suffixes) && not ("." `BS.isPrefixOf` fileName)
 165     then handleFile func [] fileName
 166     else return ()
 167 handleFile func [] fileName =
 168   X.catch
 169     (readFileRawLazy fileName >>= func (Just fileName))
 170     (\e -> printparam "exception while processing" fileName >> printparam "exception in handleFile" (e :: X.IOException))
 171 #endif
 172
 173 -- | Publish one message with our settings
 174 publishOneMsg' ::
 175      Channel
 176   -> Args
 177   -> String
 178   -> String
 179   -> Maybe RawFilePath
 180   -> BL.ByteString
 181   -> IO ()
 182 publishOneMsg' chan a exchange rkey fn content = do
 183   printparam "sending" fn
 184   (mtype, mencoding) <-
 185     if (magic a)
 186       then do
 187         let firstchunk = if BL.null content then BS.empty else head $ BL.toChunks content
 188         m <- magicOpen [MagicMimeType]
 189         magicLoadDefault m
 190         t <- BS.useAsCStringLen firstchunk (magicCString m)
 191         magicSetFlags m [MagicMimeEncoding]
 192         e <- BS.useAsCStringLen firstchunk (magicCString m)
 193         return (Just (T.pack t), Just (T.pack e))
 194       else return ((contenttype a), (contentencoding a))
 195   now <- getCurrentTime >>= return . floor . utcTimeToPOSIXSeconds
 196   publishMsg
 197     chan
 198     (T.pack $ exchange)
 199     (T.pack $ rkey)
 200     newMsg
 201       { msgBody = content
 202       , msgDeliveryMode = persistent a
 203       , msgTimestamp = Just now
 204       , msgID = msgid a
 205       , msgType = msgtype a
 206       , msgUserID = userid a
 207       , msgApplicationID = appid a
 208       , msgClusterID = clusterid a
 209       , msgContentType = mtype
 210       , msgContentEncoding = mencoding
 211       , msgReplyTo = replyto a
 212       , msgPriority = prio a
 213       , msgCorrelationID = corrid a
 214       , msgExpiration = msgexp a
 215       , msgHeaders = substheader (fnheader a) (fmap takeFileName fn) $ msgheader a
 216       } >>=
 217     printparam "sent"
 218   removeSentFileIfRequested (removeSentFile a) (moveSentFileTo a) fn
 219   where
 220     substheader ::
 221          [String] -> Maybe BS.ByteString -> Maybe FieldTable -> Maybe FieldTable
 222     substheader (s:r) (Just fname) old =
 223       substheader r (Just fname) (addheader' old s fname)
 224     substheader _ _ old = old
 225     removeSentFileIfRequested False _ _ = return ()
 226     removeSentFileIfRequested True _ Nothing = return ()
 227     removeSentFileIfRequested True Nothing (Just fname) =
 228       printparam "removing" fname >> RD.removeFile fname
 229     removeSentFileIfRequested True (Just path) (Just fname) =
 230       printparam "moving" [fname,"to",path] >>
 231       F.rename fname (replaceDirectory fname ((takeDirectory fname) </> path))
 232     addheader' :: Maybe FieldTable -> String -> BS.ByteString -> Maybe FieldTable
 233     addheader' Nothing k v =
 234       Just $ FieldTable $ M.singleton (T.pack k) (FVString v)
 235     addheader' (Just (FieldTable oldheader)) k v =
 236       Just $ FieldTable $ M.insert (T.pack k) (FVString v) oldheader