summaryrefslogtreecommitdiffstats
path: root/src/main.hs
blob: 042c8da674f4d24bcc001fffb439c1757bf9acb4 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
{-# 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 Data.Int (Int32)
import qualified Data.Text.IO as T
import Data.Maybe (fromMaybe)
import Data.Word (Word32)
import DBus
import DBus.Client
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
import System.Random
import Text.Read (readMaybe)


main :: IO ()
main =
    bracket connectSession disconnect app

app :: Client -> IO ()
app client = do
    (captureFinish, awaitCaptureFinish) <- newRelay
    (daemonJoin, awaitDaemonJoin) <- newSemaphore
    (daemonPart, awaitDaemonPart) <- newSemaphore
    (daemonStop, awaitDaemonStop) <- newSemaphore

    cGuiPath <- getEnv "FLAMESHOT_CAPTURE_PATH"

    let
        cLogTime = True
        cLogHandle = stderr

    cTimeout <-
      let
        def = 100
        err = error $ "failed to read integer from " <> var
        var = "FLAMESHOT_ONCE_TIMEOUT"
      in
        (1000*) . maybe def (fromMaybe err . readMaybe) <$> lookupEnv var

    hSetBuffering cLogHandle LineBuffering
    logToTTY <- hIsTerminalDevice cLogHandle
    (putLog0, takeLog) <- newChan

    let
        putLog1 = if cLogTime then (putLog0 =<<) . prefixTimestamp else putLog0
        putLog2 = putLog1 . showUnprintable
        putLog3 = if logToTTY then putLog2 else putLog2 . stripSGR
        putLog = putLog3

    _ <-
      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
              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 ()

    let
      runDaemon =
        P.run "flameshot" [] Nothing Nothing "" Nothing $
          logger putLog "daemon" <>
          mempty
            { P.onError = \_ _ -> daemonStop
            , P.onExit = \_ _ -> daemonStop
            }
      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 runDaemon $ \daemonThread -> do
        withTimeout_ cTimeout awaitDaemonJoin "daemon" "join" $ do
          callId <- getStdRandom random :: IO Word32
          call client (methodCall "/" "" "graphicCapture")
              { methodCallDestination = Just "org.dharkael.Flameshot"
              , methodCallInterface = Nothing
              , methodCallBody =
                [ toVariant cGuiPath
                , toVariant (0 :: Int32)
                , toVariant callId
                ]
              } >>= \case
            Right MethodReturn{} ->
              awaitCaptureFinish >>= mapM_ (copyToClipboard "image/png")

            Left err ->
              putLog $ "capture failed: " <> blessMethodError err

          killThread daemonThread

          withTimeout_ cTimeout awaitDaemonPart "daemon" "part" (return ())
          withTimeout_ cTimeout awaitDaemonStop "daemon" "stop" (return ())