initial commitv1.0.0-rc1

8 files changed, 479 insertions, 0 deletions

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 <tv@krebsco.de>
+maintainer: tv <tv@krebsco.de>
+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 <stockholm>;
+  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 ())