From 6b35bf4fcd777004657adbdcb46c507a3d1cbd1a Mon Sep 17 00:00:00 2001
From: Ayush <ayushsharma@iitb.ac.in>
Date: Tue, 26 Apr 2022 16:10:58 +0530
Subject: [PATCH] Update bluetooth_battery.py

---
 Dockerfile           |  4 ++--
 Readme.md            |  4 ++--
 bluetooth_battery.py | 32 ++++++++++++++++++++++++++++----
 3 files changed, 32 insertions(+), 8 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index bf069b1..fa08173 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -2,13 +2,13 @@
 FROM python:3.9-slim-bullseye as builder
 RUN apt update && apt install -y build-essential bluetooth libbluetooth-dev
 WORKDIR /app
-RUN pip3 wheel --no-cache-dir --no-deps --wheel-dir /app/wheel/ pybluez
+RUN pip3 wheel --no-cache-dir --no-deps --wheel-dir /app/wheel/ pybluez pydbus
 
 ### Final
 FROM python:3.9-slim-bullseye
 RUN apt update && apt install -y libbluetooth3
 WORKDIR /app
 COPY --from=builder /app/wheel /wheels
-RUN pip3 install --no-cache /wheels/*
+RUN pip3 install --no-cache /wheels/* pydbus
 COPY ./bluetooth_battery.py .
 ENTRYPOINT ["python3", "./bluetooth_battery.py"]
diff --git a/Readme.md b/Readme.md
index 0297023..59f59f8 100644
--- a/Readme.md
+++ b/Readme.md
@@ -27,7 +27,7 @@ Table of Contents
 ### There are four options:
 
 ### Option 1: Install from PyPI
-Please ensure you have the BlueZ and python libraries and header files if you are using Ubuntu/Debian based distros:
+Please ensure you have the BlueZ, pydbus and python libraries and header files if you are using Ubuntu/Debian based distros:
 ```console
 sudo apt install libbluetooth-dev python3-dev
 ```
@@ -56,7 +56,7 @@ chmod +x bluetooth_battery.py
 ./bluetooth_battery.py BT_MAC_ADDRESS_1 ...
 ```
 
-_make sure you have `python-pybluez` or `python3-pybluez` or `python3-bluez` installed on your system._
+_make sure you have `python-pybluez` or `python3-pybluez` or `python3-bluez` and `python3-pydbus` installed on your system._
 
 --------
 
diff --git a/bluetooth_battery.py b/bluetooth_battery.py
index 0005136..1f463e2 100755
--- a/bluetooth_battery.py
+++ b/bluetooth_battery.py
@@ -10,6 +10,7 @@
 
 import argparse
 import bluetooth
+import pydbus
 from typing import Optional, Union, List, Dict
 
 
@@ -148,19 +149,42 @@ def _perform_query(self) -> int:
         return result
 
 
+def return_connected_devices(mngr):
+    result={}
+    mngd_objs = mngr.GetManagedObjects()
+    for path in mngd_objs:
+        con_state = mngd_objs[path].get('org.bluez.Device1', {}).get('Connected', False)
+        if con_state:
+            addr = mngd_objs[path].get('org.bluez.Device1', {}).get('Address')
+            name = mngd_objs[path].get('org.bluez.Device1', {}).get('Name')
+            result[addr] = name
+    return result
+
+
 def main():
     """
     The starting point of the program. For each device address in the argument
     list a bluetooth socket will be opened and the battery level will be read
     and printed to stdout
     """
+    bus = pydbus.SystemBus()
+
+    adapter = bus.get('org.bluez', '/org/bluez/hci0')
+    mngr = bus.get('org.bluez', '/')
     parser = argparse.ArgumentParser(description="Get battery level from Bluetooth headsets")
-    parser.add_argument("devices", metavar="DEVICE_MAC[.PORT]", type=str, nargs="+",
+    parser.add_argument("-a","--all",action="store_true",help="Show battery of all connected devices")
+    parser.add_argument("--devices", metavar="DEVICE_MAC[.PORT]", type=str, nargs="+",
                         help="(MAC address of target)[.SPP Port]")
     args = parser.parse_args()
-    for device in args.devices:
-        query = BatteryStateQuerier(*device.split("."))
-        print("Battery level for {} is {}".format(device, str(query)))
+    if args.devices:
+        for device in args.devices:
+            query = BatteryStateQuerier(*device.split("."))
+            print("Battery level for {} is {}".format(device, str(query)))
+    if args.all:
+        result = return_connected_devices(mngr)
+        for key in result:
+            query = BatteryStateQuerier(key)
+            print("Battery level for {} is {}".format(result[key], str(query)))
 
 if __name__ == "__main__":
     main()