From fb2086171e65b61318a6ded4d7a2171c5c7cd13c Mon Sep 17 00:00:00 2001
From: VedantSawant616 <vedusawant616@gmail.com>
Date: Wed, 24 Sep 2025 23:40:23 +0530
Subject: [PATCH] feat: Add system bus logic and refactor panctl

---
 pantalaimon/bus_helper.py | 31 +++++++++++++++++++++++++++++++
 pantalaimon/daemon.py     | 21 +++++++++++++++++++++
 pantalaimon/panctl.py     | 33 +++++++++++++++++++++++++++++----
 3 files changed, 81 insertions(+), 4 deletions(-)
 create mode 100644 pantalaimon/bus_helper.py

diff --git a/pantalaimon/bus_helper.py b/pantalaimon/bus_helper.py
new file mode 100644
index 0000000..4cda92b
--- /dev/null
+++ b/pantalaimon/bus_helper.py
@@ -0,0 +1,31 @@
+# This new file centralizes the logic for D-Bus selection.
+
+import os
+import dbus
+
+from dbus.exceptions import DBusException
+
+def get_bus(bus_choice):
+    """
+    Returns a D-Bus connection based on the user's choice.
+    Raises RuntimeError if no connection can be made.
+    """
+    if bus_choice == "system":
+        return dbus.SystemBus()
+    elif bus_choice == "session":
+        return dbus.SessionBus()
+    elif bus_choice == "auto":
+        # First, try to connect to the session bus. This is the default
+        # behavior for per-user installs.
+        try:
+            return dbus.SessionBus()
+        except DBusException:
+            pass # Fallback to system bus if session bus fails
+
+        # If session bus fails, try system bus.
+        try:
+            return dbus.SystemBus()
+        except dbus.DBusException:
+            raise RuntimeError("Could not find pantalaimon on session or system bus")
+    else:
+        raise ValueError(f"Invalid bus choice: {bus_choice}")
\ No newline at end of file
diff --git a/pantalaimon/daemon.py b/pantalaimon/daemon.py
index 6d47b36..4ed32c6 100755
--- a/pantalaimon/daemon.py
+++ b/pantalaimon/daemon.py
@@ -12,6 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
+import argparse
+from pantalaimon.bus_helper import get_bus
+
 import asyncio
 import json
 import os
@@ -1394,3 +1398,20 @@ async def shutdown(self, _):
         if self.default_session:
             await self.default_session.close()
             self.default_session = None
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Pantalaimon Daemon")
+    parser.add_argument(
+        "--bus",
+        choices=["session", "system", "auto"],
+        default="auto",
+        help="Choose which D-Bus bus to connect to (default: auto)",
+    )
+    args = parser.parse_args()
+
+    # Use the helper to get the bus
+    bus = get_bus(args.bus)
+
+    # TODO: Place daemon startup logic here
+    print(f"Daemon started with {args.bus} bus: {bus}")
diff --git a/pantalaimon/panctl.py b/pantalaimon/panctl.py
index f37ab88..0cd51d8 100644
--- a/pantalaimon/panctl.py
+++ b/pantalaimon/panctl.py
@@ -35,7 +35,9 @@
 from prompt_toolkit.completion import Completer, Completion, PathCompleter
 from prompt_toolkit.document import Document
 from prompt_toolkit.patch_stdout import patch_stdout
-from pydbus import SessionBus
+from dbus.exceptions import DBusException
+from pantalaimon.bus_helper import get_bus
+
 
 PTK2 = ptk_version.startswith("2.")
 
@@ -404,8 +406,31 @@ class PanCtl:
 
     commands = list(command_help.keys())
 
-    def __attrs_post_init__(self):
-        self.bus = SessionBus()
+   def __attrs_post_init__(self):
+        try:
+            self.bus = get_bus("auto")
+        except RuntimeError as e:
+            raise DBusException(f"Could not connect to D-Bus: {e}")
+
+        self.pan_bus = self.bus.get("org.pantalaimon1")
+
+        self.ctl = self.pan_bus["org.pantalaimon1.control"]
+        self.devices = self.pan_bus["org.pantalaimon1.devices"]
+
+        self.own_message_ids = []
+
+        self.ctl.Response.connect(self.show_response)
+        self.ctl.UnverifiedDevices.connect(self.unverified_devices)
+
+        self.completer = PanCompleter(self.commands, self.ctl, self.devices)
+
+        self.devices.VerificationInvite.connect(self.show_sas_invite)
+        self.devices.VerificationString.connect(self.show_sas)
+        self.devices.VerificationDone.connect(self.sas_done)
+
+        self.devices.KeyRequest.connect(self.show_key_request)
+        self.devices.KeyRequestCancel.connect(self.show_key_request_cancel)
+
         self.pan_bus = self.bus.get("org.pantalaimon1")
 
         self.ctl = self.pan_bus["org.pantalaimon1.control"]
@@ -703,7 +728,7 @@ def main():
 
     try:
         panctl = PanCtl()
-    except GLib.Error as e:
+    except (DBusException, RuntimeError) as e:
         print(f"Error, {e}")
         sys.exit(-1)