From fb5636483871fbafe9b286b377c339c8ddf8b4f8 Mon Sep 17 00:00:00 2001 From: tv Date: Thu, 7 Feb 2019 18:42:36 +0100 Subject: initial commit --- flameshot-once.cabal | 24 +++++ shell.nix | 35 +++++++ src/Control/Concurrent/Extended.hs | 24 +++++ src/Control/Monad/Extended.hs | 14 +++ src/Data/Text/Extended.hs | 12 +++ src/Flameshot/Internal.hs | 96 +++++++++++++++++++ src/Flameshot/Internal/Process.hs | 89 ++++++++++++++++++ src/main.hs | 185 +++++++++++++++++++++++++++++++++++++ 8 files changed, 479 insertions(+) create mode 100644 flameshot-once.cabal create mode 100644 shell.nix create mode 100644 src/Control/Concurrent/Extended.hs create mode 100644 src/Control/Monad/Extended.hs create mode 100644 src/Data/Text/Extended.hs create mode 100644 src/Flameshot/Internal.hs create mode 100644 src/Flameshot/Internal/Process.hs create mode 100644 src/main.hs diff --git a/flameshot-once.cabal b/flameshot-once.cabal new file mode 100644 index 0000000..c38a65b --- /dev/null +++ b/flameshot-once.cabal @@ -0,0 +1,24 @@ +name: flameshot-once +version: 1.0.0 +license: MIT +author: tv +maintainer: tv +build-type: Simple +cabal-version: >=1.10 + +executable flameshot-once + build-depends: + async, + base, + blessings, + dbus, + process, + iso8601-time, + text, + time, + unagi-chan, + unix + default-language: Haskell2010 + ghc-options: -O2 -Wall -threaded + hs-source-dirs: src + main-is: main.hs diff --git a/shell.nix b/shell.nix new file mode 100644 index 0000000..4d7e83a --- /dev/null +++ b/shell.nix @@ -0,0 +1,35 @@ +{ compiler ? "default" }: let + + stockholm = import ; + inherit (stockholm.systems.${lib.krops.getHostName}) config pkgs; + inherit (stockholm) lib; + + haskellPackages = + if compiler == "default" + then pkgs.haskellPackages + else pkgs.haskell.packages.${compiler}; + + drv = haskellPackages.callPackage (import ./.) {}; + +in + + lib.overrideDerivation drv.env (oldAttrs: { + + buildInputs = [ + pkgs.flameshot + ]; + + shellHook = '' + pkg_name=${lib.escapeShellArg (builtins.baseNameOf (toString ./.))} + + WORKDIR=${lib.escapeShellArg (toString ./.)} + CACHEDIR=$HOME/tmp/$pkg_name + HISTFILE=$CACHEDIR/bash_history + + mkdir -p "$CACHEDIR" + + export SHELL=/run/current-system/sw/bin/bash + + cd "$WORKDIR" + ''; + }) 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 ()) -- cgit v1.2.3