From c0367b2c759dea366b09a1fd373630f6951cd924 Mon Sep 17 00:00:00 2001
From: Bernd Schoolmann <mail@quexten.com>
Date: Sat, 23 Dec 2023 08:37:17 +0100
Subject: [PATCH] Add lock on screensaver

---
 agent/processsecurity/unimplemented.go |  4 ++
 agent/processsecurity/unix.go          | 52 +++++++++++++++++++++++++-
 agent/unixsocketagent.go               |  9 +++++
 agent/virtualagent.go                  |  9 +++++
 4 files changed, 73 insertions(+), 1 deletion(-)

diff --git a/agent/processsecurity/unimplemented.go b/agent/processsecurity/unimplemented.go
index 0774dea..9c5dafa 100644
--- a/agent/processsecurity/unimplemented.go
+++ b/agent/processsecurity/unimplemented.go
@@ -6,3 +6,7 @@ func DisableDumpable() error {
 	// no additional dumping protection
 	return nil
 }
+
+func MonitorLocks(onlock func) error {
+	return nil
+}
diff --git a/agent/processsecurity/unix.go b/agent/processsecurity/unix.go
index f3579fc..d14704d 100644
--- a/agent/processsecurity/unix.go
+++ b/agent/processsecurity/unix.go
@@ -2,8 +2,58 @@
 
 package processsecurity
 
-import "golang.org/x/sys/unix"
+import (
+	"fmt"
+
+	"github.com/godbus/dbus/v5"
+	"golang.org/x/sys/unix"
+)
 
 func DisableDumpable() error {
 	return unix.Prctl(unix.PR_SET_DUMPABLE, 0, 0, 0, 0)
 }
+
+func MonitorLocks(onlock func()) error {
+	bus, err := dbus.SessionBus()
+	if err != nil {
+		return err
+	}
+	err = bus.AddMatchSignal(dbus.WithMatchInterface("org.gnome.ScreenSaver"))
+	if err != nil {
+		return err
+	}
+	err = bus.AddMatchSignal(dbus.WithMatchMember("org.freedesktop.ScreenSaver"))
+	if err != nil {
+		return err
+	}
+
+	signals := make(chan *dbus.Signal, 10)
+	bus.Signal(signals)
+	for {
+		select {
+		case message := <-signals:
+			fmt.Println("Message:", message)
+			fmt.Println("name ", message.Name)
+			if message.Name == "org.gnome.ScreenSaver.ActiveChanged" {
+				if len(message.Body) == 0 {
+					continue
+				}
+				locked, err := message.Body[0].(bool)
+				if err || locked {
+					onlock()
+				}
+			}
+			if message.Name == "org.freedesktop.ScreenSaver.ActiveChanged" {
+				if len(message.Body) == 0 {
+					continue
+				}
+				locked, err := message.Body[0].(bool)
+				if err || locked {
+					onlock()
+				}
+			}
+		}
+	}
+
+	return nil
+}
diff --git a/agent/unixsocketagent.go b/agent/unixsocketagent.go
index b4e48e1..76a7303 100644
--- a/agent/unixsocketagent.go
+++ b/agent/unixsocketagent.go
@@ -153,6 +153,15 @@ func StartUnixAgent(path string, runtimeConfig config.RuntimeConfig) error {
 	}
 
 	processsecurity.DisableDumpable()
+	err = processsecurity.MonitorLocks(func() {
+		cfg.Lock()
+		vault.Clear()
+		vault.Keyring.Lock()
+	})
+	if err != nil {
+		log.Warn("Could not monitor screensaver: %s", err.Error())
+	}
+
 	if !runtimeConfig.WebsocketDisabled {
 		go bitwarden.RunWebsocketDaemon(ctx, vault, &cfg)
 	}
diff --git a/agent/virtualagent.go b/agent/virtualagent.go
index f8b3e33..156db4a 100644
--- a/agent/virtualagent.go
+++ b/agent/virtualagent.go
@@ -116,6 +116,15 @@ func StartVirtualAgent(runtimeConfig config.RuntimeConfig) (chan []byte, chan []
 		}
 	}
 	processsecurity.DisableDumpable()
+	err = processsecurity.MonitorLocks(func() {
+		cfg.Lock()
+		vault.Clear()
+		vault.Keyring.Lock()
+	})
+	if err != nil {
+		log.Warn("Could not monitor screensaver: %s", err.Error())
+	}
+
 	go func() {
 		for {
 			time.Sleep(TokenRefreshInterval)