From 20acd8762576221847daea25b00100bc33242e7e Mon Sep 17 00:00:00 2001
From: Lenoch <81lennoch@gmail.com>
Date: Mon, 16 Dec 2024 16:34:00 +0100
Subject: [PATCH] Add support for locking session on Linux and refactor
 functions

- Added support for locking session on Linux with various methods
- Moved lock session function to a dedicated lock_session.py file
- Updated commands.py to use the new lock_session function
---
 app/buttons/commands.py            |  2 +-
 app/buttons/system/__init__.py     |  1 +
 app/buttons/system/lock_session.py | 30 ++++++++++++++++++++++++++++++
 3 files changed, 32 insertions(+), 1 deletion(-)
 create mode 100644 app/buttons/system/lock_session.py

diff --git a/app/buttons/commands.py b/app/buttons/commands.py
index a7a9ad9..2632920 100644
--- a/app/buttons/commands.py
+++ b/app/buttons/commands.py
@@ -72,7 +72,7 @@ def handle_command(message=None):
         return soundboard.playsound(*soundboard.get_params(message))
 
     elif message.startswith("/locksession"):
-        subprocess.Popen("Rundll32.exe user32.dll,LockWorkStation", shell=True)
+        system.lock_session()
 
     elif message.startswith("/screensaversettings"):
         subprocess.Popen(
diff --git a/app/buttons/system/__init__.py b/app/buttons/system/__init__.py
index c72b7c3..4281f36 100644
--- a/app/buttons/system/__init__.py
+++ b/app/buttons/system/__init__.py
@@ -1,2 +1,3 @@
 from .command_handler import handle_command
 from .screensaver import handle_command as screensaver
+from .lock_session import lock_session
diff --git a/app/buttons/system/lock_session.py b/app/buttons/system/lock_session.py
new file mode 100644
index 0000000..a14e8af
--- /dev/null
+++ b/app/buttons/system/lock_session.py
@@ -0,0 +1,30 @@
+from app.utils.platform import is_windows, is_linux, is_mac
+
+import subprocess
+from app.utils.logger import log
+
+
+def lock_session():
+    """Lock the user's session."""
+    if is_windows:
+        # Lock workstation on Windows
+        subprocess.Popen("Rundll32.exe user32.dll,LockWorkStation", shell=True)
+        
+    elif is_linux:
+        # Try various Linux locking methods
+        if subprocess.run("dm-tool lock", shell=True).returncode == 0:
+            return
+        elif subprocess.run("xdg-screensaver lock", shell=True).returncode == 0:
+            return
+        elif subprocess.run("gnome-screensaver-command -l", shell=True).returncode == 0:
+            return
+        elif subprocess.run("loginctl lock-session", shell=True).returncode == 0:
+            return
+        elif subprocess.run("swaylock", shell=True).returncode == 0:
+            return
+        else:
+            log.error("Unable to lock the session.")
+            raise RuntimeError("Unable to lock the session.")
+
+    else:
+        raise NotImplementedError("Screensaver is not implemented for this platform.")