fix: Use grim to take screenshots on wayland (#6)

README.md

@@ -19,6 +19,9 @@ TLock is an open-source tool to store and manage your authentication tokens secu
 >[!NOTE]
 >For showing the provider's icon, you must have Nerd Fonts installed
 
+>[!NOTE]
+>To make screenshot feature work on wayland, you need to install grim
+
 ## ⬇️ Installation
 
 - **Arch Linux** (with AUR helper, like yay)

tlock-internal/utils/read_screen.go

@@ -0,0 +1,110 @@
+package utils
+
+import (
+	"errors"
+	"fmt"
+	"image"
+	"os"
+	"os/exec"
+
+	_ "image/png"
+
+	"github.com/kbinani/screenshot"
+	"github.com/makiuchi-d/gozxing"
+	"github.com/makiuchi-d/gozxing/qrcode"
+)
+
+// Loads an image from file
+func getImageFromFilePath(filePath string) (image.Image, error) {
+	// Read file path
+	f, err := os.Open(filePath)
+	if err != nil {
+		return nil, err
+	}
+
+	// Close on scope end
+	defer f.Close()
+
+	// Decode
+	image, _, err := image.Decode(f)
+
+	// Return
+	return image, err
+}
+
+// Returns if the current session is wayland
+func isWayland() bool {
+	return os.Getenv("XDG_SESSION_TYPE") == "wayland"
+}
+
+// Reads the QRCode from the image
+func readFromImage(image image.Image) (*string, error) {
+	// Create bitmap
+	bmp, err := gozxing.NewBinaryBitmapFromImage(image)
+	if err != nil {
+		return nil, err
+	}
+
+	// Create reader
+	qrReader := qrcode.NewQRCodeReader()
+
+	// Decode
+	if result, err := qrReader.Decode(bmp, nil); err == nil {
+		// Get fetched URI
+		uri := result.String()
+
+		// Return
+		return &uri, nil
+	}
+
+	// No token found
+	return nil, TOKEN_NOT_FOUND_ERR
+
+}
+
+// Error when no token on the screen is found
+var TOKEN_NOT_FOUND_ERR = errors.New("No token found on the screen")
+
+// Reads QRCode from screen and returns the found data (not for wayland)
+func ReadTokenFromScreenNoWayland() (*string, error) {
+	// Capture rect
+	image, err := screenshot.CaptureRect(screenshot.GetDisplayBounds(0))
+	if err != nil {
+		return nil, err
+	}
+
+	// Return
+	return readFromImage(image)
+}
+
+// Reads QRCode form screen and returns the found data (only for wayland)
+func ReadTokenFromScreenWaylandOnly() (*string, error) {
+	// Out path
+	out := "/tmp/tlock_screenshot.png"
+
+	// Make grim command
+	_, err := exec.Command("grim", out).Output()
+	if err != nil {
+		return nil, errors.New(fmt.Sprintf("Cannot run grim command: %s", err))
+	}
+
+	// Load screenshot
+	image, err := getImageFromFilePath(out)
+	if err != nil {
+		return nil, err
+	}
+
+	// Read
+	return readFromImage(image)
+}
+
+// Reads QRCode from screen based on the session type
+func ReadTokenFromScreen() (*string, error) {
+	// Check if it is wayland
+	if isWayland() {
+		return ReadTokenFromScreenWaylandOnly()
+	}
+
+	// Use normal function
+	return ReadTokenFromScreenNoWayland()
+}

tlock/models/dashboard/tokens/from_screen.go

@@ -2,7 +2,6 @@ package tokens
 
 import (
 	"fmt"
-	"strings"
 	"time"
 
 	"github.com/charmbracelet/bubbles/key"
@@ -12,11 +11,9 @@ import (
 	"github.com/eklairs/tlock/tlock-internal/components"
 	tlockmessages "github.com/eklairs/tlock/tlock-internal/messages"
 	"github.com/eklairs/tlock/tlock-internal/modelmanager"
+	"github.com/eklairs/tlock/tlock-internal/utils"
 	tlockvault "github.com/eklairs/tlock/tlock-vault"
 	tlockstyles "github.com/eklairs/tlock/tlock/styles"
-	"github.com/kbinani/screenshot"
-	"github.com/makiuchi-d/gozxing"
-	"github.com/makiuchi-d/gozxing/qrcode"
 	"github.com/pquerna/otp"
 )
 
@@ -190,29 +187,30 @@ func (screen TokenFromScreen) Update(msg tea.Msg, manager *modelmanager.ModelMan
 			cmds = append(cmds, screen.spinner.Tick)
 
 			go func() {
-				if image, err := screenshot.CaptureRect(screenshot.GetDisplayBounds(0)); err == nil {
-					if bmp, err := gozxing.NewBinaryBitmapFromImage(image); err == nil {
-						qrReader := qrcode.NewQRCodeReader()
-
-						if result, err := qrReader.Decode(bmp, nil); err == nil {
-							uri := result.String()
-
-							// Try to parse
-							if key, err := otp.NewKeyFromURL(uri); err == nil {
-								// Run validator
-								_, err := screen.vault.ValidateToken(key.Secret())
-
-								// Send
-								dataFromScreenChan <- &dataFromScreen{
-									Uri: &uri,
-									Err: err,
-								}
-							}
+				// Read data from the screen
+				data, err := utils.ReadTokenFromScreen()
+
+				// Prepare data
+				dataScreen := dataFromScreen{
+					Uri: data,
+					Err: err,
+				}
+
+				// Try to parse
+				if data != nil {
+					if key, err := otp.NewKeyFromURL(*data); err == nil {
+						// Run validator
+						_, err := screen.vault.ValidateToken(key.Secret())
+
+						// Update error message
+						if dataScreen.Err == nil {
+							dataScreen.Err = err
 						}
 					}
 				}
 
-				dataFromScreenChan <- nil
+				// Send
+				dataFromScreenChan <- &dataScreen
 			}()
 
 		case key.Matches(msgType, confirmScreenKeys.Retake):
@@ -228,6 +226,7 @@ func (screen TokenFromScreen) Update(msg tea.Msg, manager *modelmanager.ModelMan
 			// Add the token
 			if screen.token != nil && screen.token.Err == nil {
 				// Add token
+				// We can ignore validations because we have already pre-checked it
 				screen.vault.AddToken(screen.folder.Name, *screen.token.Uri)
 
 				// Require refresh of folders and tokens list
@@ -282,9 +281,7 @@ func (screen TokenFromScreen) View() string {
 		}
 
 		// If the token is null, show the message
-		if screen.token == nil {
-			items = append(items, tlockstyles.Styles.Error.Render("Did not find any token!"))
-		} else {
+		if screen.token != nil {
 			// Try to parse the otp value
 			if screen.token.Err == nil {
 				key, err := otp.NewKeyFromURL(*screen.token.Uri)
@@ -312,8 +309,7 @@ func (screen TokenFromScreen) View() string {
 			} else {
 				items = append(items, lipgloss.JoinHorizontal(
 					lipgloss.Center,
-					tlockstyles.Styles.Base.Render("Found a token from screen, but "),
-					tlockstyles.Styles.Error.Render(strings.ToLower(screen.token.Err.Error())),
+					tlockstyles.Styles.Error.Render(screen.token.Err.Error()),
 				))
 			}
 		}