summaryrefslogtreecommitdiffstats
path: root/src
diff options
context:
space:
mode:
Diffstat (limited to 'src')
-rw-r--r--src/Control/Concurrent/Extended.hs24
-rw-r--r--src/Control/Monad/Extended.hs14
-rw-r--r--src/Data/Text/Extended.hs12
-rw-r--r--src/Flameshot/Internal.hs96
-rw-r--r--src/Flameshot/Internal/Process.hs89
-rw-r--r--src/main.hs185
6 files changed, 420 insertions, 0 deletions
diff --git a/src/Control/Concurrent/Extended.hs b/src/Control/Concurrent/Extended.hs
new file mode 100644
index 0000000..933e3a6
--- /dev/null
+++ b/src/Control/Concurrent/Extended.hs
@@ -0,0 +1,24 @@
+module Control.Concurrent.Extended
+ ( module Exports
+ , newChan
+ , newRef
+ , newRelay
+ , newSemaphore
+ ) where
+
+import Control.Arrow
+import Control.Concurrent as Exports hiding (newChan,readChan,writeChan)
+import qualified Control.Concurrent.Chan.Unagi as U
+import Data.IORef
+
+newChan :: IO (a -> IO (), IO a)
+newChan = (U.writeChan *** U.readChan) <$> U.newChan
+
+newRef :: a -> IO (a -> IO (), IO a)
+newRef v0 = (atomicWriteIORef &&& readIORef) <$> newIORef v0
+
+newRelay :: IO (a -> IO (), IO a)
+newRelay = (putMVar &&& takeMVar) <$> newEmptyMVar
+
+newSemaphore :: IO (IO (), IO ())
+newSemaphore = first ($()) <$> newRelay
diff --git a/src/Control/Monad/Extended.hs b/src/Control/Monad/Extended.hs
new file mode 100644
index 0000000..d91b12c
--- /dev/null
+++ b/src/Control/Monad/Extended.hs
@@ -0,0 +1,14 @@
+module Control.Monad.Extended
+ ( module Control.Monad
+ , unlessM_
+ , untilM_
+ ) where
+
+import Control.Monad
+
+
+unlessM_ :: Monad m => m Bool -> m () -> m ()
+unlessM_ p f = p >>= flip unless f
+
+untilM_ :: Monad m => m Bool -> m a -> m ()
+untilM_ p f = unlessM_ p (f >> untilM_ p f)
diff --git a/src/Data/Text/Extended.hs b/src/Data/Text/Extended.hs
new file mode 100644
index 0000000..70eef63
--- /dev/null
+++ b/src/Data/Text/Extended.hs
@@ -0,0 +1,12 @@
+{-# OPTIONS_GHC -fno-warn-orphans #-}
+module Data.Text.Extended
+ ( module Data.Text
+ , show
+ ) where
+
+import Data.Text
+import Prelude hiding (show)
+import qualified Prelude
+
+show :: Show a => a -> Text
+show = pack . Prelude.show
diff --git a/src/Flameshot/Internal.hs b/src/Flameshot/Internal.hs
new file mode 100644
index 0000000..28b33cf
--- /dev/null
+++ b/src/Flameshot/Internal.hs
@@ -0,0 +1,96 @@
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE MultiWayIf #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE RecordWildCards #-}
+
+module Flameshot.Internal where
+
+import Blessings.Text
+import Control.Concurrent.Async (race)
+import Control.Concurrent.Extended
+import Control.Exception
+import qualified Data.Char as C
+import Data.Function (on)
+import Data.Time.Clock.System
+import Data.Time.ISO8601
+import DBus
+import DBus.Socket
+import Data.Text (Text)
+import qualified Data.Text.Extended as T
+
+
+blessBusName :: BusName -> Blessings Text
+blessBusName = Plain . T.pack . formatBusName
+
+blessMemberName :: MemberName -> Blessings Text
+blessMemberName = Plain . T.pack . formatMemberName
+
+blessShow :: Show a => a -> Blessings Text
+blessShow = Plain . T.show
+
+blessTime :: Int -> Blessings Text
+blessTime = Plain . (<>"μs") . T.show
+
+red :: Blessings Text -> Blessings Text
+red = SGR [31]
+
+
+
+withThread :: IO () -> (ThreadId -> IO a) -> IO a
+withThread tf wtf =
+ bracket (forkIO tf) killThread wtf
+
+withThread_ :: IO () -> IO a -> IO a
+withThread_ tf wtf =
+ bracket (forkIO tf) killThread (const wtf)
+
+
+timeout :: Int -> IO a -> IO (Either Int a)
+timeout time io =
+ race (threadDelay time) io >>= \case
+ Right x -> return (Right x)
+ Left () -> return (Left time)
+
+
+dbusInterface :: InterfaceName
+dbusInterface = "org.freedesktop.DBus"
+
+withSocket :: Address -> (Socket -> IO a) -> IO a
+withSocket addr = bracket (open addr) close
+
+
+getTimestamp :: IO String
+getTimestamp =
+ formatISO8601Micros . systemToUTCTime <$> getSystemTime
+
+prefixTimestamp :: Blessings Text -> IO (Blessings Text)
+prefixTimestamp s =
+ (<> s) . (<>" ") . SGR [38,5,239] . Plain . T.pack <$> getTimestamp
+
+showUnprintable :: Blessings Text -> Blessings Text
+showUnprintable =
+ fmap' showU
+ where
+ showU :: Text -> Blessings Text
+ showU =
+ mconcat
+ . map (either Plain (hi . Plain . showLitChars))
+ . toEither (not . C.isPrint)
+
+ -- like Blessings' fmap, but don't wrap the Plain case in another Plain
+ fmap' :: (Text -> Blessings Text) -> Blessings Text -> Blessings Text
+ fmap' f = \case
+ Append t1 t2 -> Append (fmap' f t1) (fmap' f t2)
+ Plain s -> f s
+ SGR pm t -> SGR pm (fmap' f t)
+ Empty -> Empty
+
+ hi = SGR [38,5,79]
+
+ showLitChars :: Text -> Text
+ showLitChars = T.concatMap (T.pack . flip C.showLitChar "")
+
+ toEither :: (Char -> Bool) -> Text -> [Either Text Text]
+ toEither p =
+ map (\s -> if p (T.head s) then Right s else Left s)
+ . T.groupBy ((==) `on` p)
diff --git a/src/Flameshot/Internal/Process.hs b/src/Flameshot/Internal/Process.hs
new file mode 100644
index 0000000..c435d48
--- /dev/null
+++ b/src/Flameshot/Internal/Process.hs
@@ -0,0 +1,89 @@
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE MultiWayIf #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE RecordWildCards #-}
+module Flameshot.Internal.Process (run,Callbacks(..)) where
+
+import Control.Concurrent (forkIO,threadDelay)
+import Control.Concurrent.Async (race)
+import Control.Exception
+import Control.Monad.Extended (untilM_,unless)
+import Data.Text (Text)
+import qualified Data.Text.IO as T
+import System.Exit
+import System.IO (BufferMode(LineBuffering),hSetBuffering)
+import System.IO (Handle,hClose,hPutStr,hIsEOF)
+import System.IO.Error (catchIOError,isDoesNotExistError)
+import System.Process
+import System.Posix.Process (getProcessGroupIDOf)
+import System.Posix.Signals (Signal,signalProcessGroup,killProcess)
+import System.Posix.Types (ProcessGroupID)
+
+data Callbacks = Callbacks
+ { onOutLine :: Pid -> Text -> IO ()
+ , onErrLine :: Pid -> Text -> IO ()
+ , onError :: Pid -> SomeException -> IO ()
+ , onExit :: Pid -> ExitCode -> IO ()
+ , onStart :: Pid -> IO ()
+ }
+
+run :: FilePath
+ -> [String]
+ -> Maybe FilePath
+ -> Maybe [(String, String)]
+ -> String
+ -> Maybe Int
+ -> Callbacks
+ -> IO ()
+run path args cwd env input hTimeout Callbacks{..} =
+ f `catch` onError (-1)
+ where
+ f = withCreateProcess p $ \(Just inh) (Just outh) (Just errh) ph -> do
+ Just pid <- getPid ph
+ pgid <- getProcessGroupIDOf pid
+
+ onStart pid
+
+ mapM_ forkIO [
+ hPutStr inh input `finally` hClose inh,
+ hWithLines outh (onOutLine pid),
+ hWithLines errh (onErrLine pid)
+ ]
+
+ case hTimeout of
+ Just time ->
+ race (threadDelay time) (waitForProcess ph) >>= \case
+ Left () -> onError pid (SomeException (ErrorCall "timeout"))
+ Right code -> onExit pid code
+ Nothing ->
+ waitForProcess ph >>= onExit pid
+
+ killProcessGroup pgid
+
+ p = (proc path args)
+ { cwd = cwd
+ , env = env
+ , std_in = CreatePipe
+ , std_out = CreatePipe
+ , std_err = CreatePipe
+ , close_fds = True
+ , create_group = True
+ , new_session = True
+ }
+
+
+
+killProcessGroup :: ProcessGroupID -> IO ()
+killProcessGroup = signalProcessGroup' killProcess
+
+signalProcessGroup' :: Signal -> ProcessGroupID -> IO ()
+signalProcessGroup' sig pgid =
+ catchIOError
+ (signalProcessGroup sig pgid)
+ (\e -> unless (isDoesNotExistError e) $ ioError e)
+
+hWithLines :: Handle -> (Text -> IO ()) -> IO ()
+hWithLines h f = do
+ hSetBuffering h LineBuffering
+ untilM_ (hIsEOF h) (T.hGetLine h >>= f) `finally` hClose h
diff --git a/src/main.hs b/src/main.hs
new file mode 100644
index 0000000..d45404f
--- /dev/null
+++ b/src/main.hs
@@ -0,0 +1,185 @@
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE MultiWayIf #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE RecordWildCards #-}
+
+module Main (main) where
+
+import Blessings.Text
+import Control.Concurrent.Extended
+import Control.Exception
+import Control.Monad
+import qualified Data.Text.Extended as T
+import qualified Data.Text.IO as T
+import DBus
+import DBus.Client
+import DBus.Socket
+import DBus.Internal.Message
+import DBus.Internal.Types
+import Flameshot.Internal
+import qualified Flameshot.Internal.Process as P
+--import System.Exit
+import System.Environment
+import System.IO
+
+
+main :: IO ()
+main =
+ getSessionAddress >>= \case
+ Just addr ->
+ withSocket addr $ \sock -> do
+ send sock (methodCall dbusPath dbusInterface "Hello")
+ { methodCallDestination = Just "org.freedesktop.DBus"
+ } (\_ser -> return ())
+ _ <- receive sock
+ bracket connectSession disconnect (app sock)
+ Nothing ->
+ error
+ (mconcat
+ [ "DBus session bus address not found; "
+ , "is DBUS_SESSION_BUS_ADDRESS set properly?"
+ ])
+
+app :: Socket -> Client -> IO ()
+app sock client = do
+ (captureFinish, awaitCaptureFinish) <- newRelay
+ (captureBegin, awaitCaptureBegin) <- newSemaphore
+ (daemonJoin, awaitDaemonJoin) <- newRelay
+ (daemonPart, awaitDaemonPart) <- newSemaphore
+ (daemonStop, awaitDaemonStop) <- newSemaphore
+
+ cGuiPath <- getEnv "FLAMESHOT_CAPTURE_PATH"
+ let
+ cLogTime = True
+ cLogHandle = stderr
+
+ hSetBuffering cLogHandle LineBuffering
+ logToTTY <- hIsTerminalDevice cLogHandle
+ (putLog, takeLog0) <- newChan
+
+ let
+ takeLog1 = if cLogTime then takeLog0 >>= prefixTimestamp else takeLog0
+ takeLog2 = showUnprintable <$> takeLog1
+ takeLog3 = if logToTTY then takeLog2 else stripSGR <$> takeLog2
+ takeLog = takeLog3
+
+ _ <-
+ addMatch client matchAny $ \case
+ Signal
+ { signalPath = "/org/freedesktop/DBus"
+ , signalInterface = "org.freedesktop.DBus"
+ , signalMember = "NameOwnerChanged"
+ , signalSender = Just "org.freedesktop.DBus"
+ , signalDestination = Nothing
+ , signalBody =
+ [ Variant (ValueAtom (AtomText "org.dharkael.Flameshot"))
+ , Variant (ValueAtom (AtomText _old_owner))
+ , Variant (ValueAtom (AtomText new_owner))
+ ]
+ }
+ ->
+ if new_owner /= ""
+ then daemonJoin (busName_ (T.unpack new_owner))
+ else daemonPart
+
+ Signal
+ { signalPath = "/"
+ , signalInterface = "org.dharkael.Flameshot"
+ , signalDestination = Nothing
+ , signalBody = Variant (ValueAtom (AtomWord32 _callId)) : rest
+ , ..
+ }
+ ->
+ case signalMember of
+ "captureTaken" -> do
+ captureFinish (Just rawImage)
+
+ "captureFailed" -> do
+ captureFinish Nothing
+
+ _other -> do
+ captureFinish Nothing
+
+ where
+ [Variant (ValueBytes rawImage)] = rest
+
+ _sig ->
+ return ()
+
+ send sock (methodCall dbusPath dbusInterface "AddMatch")
+ { methodCallDestination = Just "org.freedesktop.DBus"
+ , methodCallBody =
+ [ toVariant (T.intercalate ","
+ [ "type='method_call'"
+ , "eavesdrop=true"
+ , "path=/"
+ , "destination=org.dharkael.Flameshot"
+ , "member=graphicCapture"
+ ] )
+ ]
+ } (\_ser -> return ())
+
+ let
+ runDaemon =
+ P.run "flameshot" [] Nothing Nothing "" Nothing P.Callbacks
+ { P.onOutLine = \_ s -> putLog $ "daemon stdout: " <> Plain s
+ , P.onErrLine = \_ s -> putLog $ "daemon stderr: " <> red (Plain s)
+ , P.onError = \_ _err -> do
+ daemonStop
+
+ , P.onExit = \_ _status -> do
+ daemonStop
+
+ , P.onStart = \_ -> return ()
+ }
+ runGui =
+ P.run "flameshot" ["gui", "-p", cGuiPath] Nothing Nothing "" Nothing P.Callbacks
+ { P.onOutLine = \_ s -> putLog $ "gui stdout: " <> Plain s
+ , P.onErrLine = \_ s -> putLog $ "gui stderr: " <> red (Plain s)
+ , P.onError = \_ _err -> return ()
+ , P.onExit = \_ _ -> return ()
+ , P.onStart = \_ -> return ()
+ }
+ runReceiver =
+ forever $ receive sock >>= \case
+ ReceivedMethodCall _ser MethodCall
+ { methodCallSender = Just _sender
+ , methodCallBody =
+ [ Variant (ValueAtom (AtomText _path))
+ , Variant (ValueAtom (AtomInt32 _delay))
+ , Variant (ValueAtom (AtomWord32 _callId))
+ ]
+ }
+ -> do
+ captureBegin
+
+ _received ->
+ return ()
+
+ runLogPrinter =
+ forever $ takeLog >>= T.hPutStrLn stderr . pp
+
+ let
+ withTimeout time f s v io =
+ timeout time f >>= \case
+ Right x ->
+ io x
+ Left _time ->
+ putLog (s <> " " <> v <> " timeout after " <> blessTime time)
+
+ withTimeout_ time f s v io =
+ withTimeout time f s v (const io)
+
+ withThread_ runLogPrinter $ do
+ withThread_ runReceiver $ do
+ withThread runDaemon $ \daemonThread -> do
+ withTimeout 100000 awaitDaemonJoin "daemon" "join" $ \_daemonBusName -> do
+ withThread_ runGui $ do
+ withTimeout_ 100000 awaitCaptureBegin "gui" "join" $ do
+ _maybeRawImage <- awaitCaptureFinish
+ return ()
+
+ killThread daemonThread
+
+ withTimeout_ 100000 awaitDaemonPart "daemon" "part" (return ())
+ withTimeout_ 100000 awaitDaemonStop "daemon" "stop" (return ())
generated by cgit v1.2.3 (git 2.25.1) at 2024-11-21 21:54:22 +0000