From 98cad3d1de83b8fabeccdaea10dacf3ec1d644bc Mon Sep 17 00:00:00 2001
From: Ivan Belokobylskiy <belokobylskij@gmail.com>
Date: Sat, 13 May 2023 23:45:27 +0200
Subject: [PATCH] voltage_bm2: initial support

---
 ble2mqtt/devices/__init__.py    |  1 +
 ble2mqtt/devices/voltage_bm2.py | 56 +++++++++++++++++++++++++++++++++
 requirements.txt                |  1 +
 setup.py                        |  3 ++
 4 files changed, 61 insertions(+)
 create mode 100644 ble2mqtt/devices/voltage_bm2.py

diff --git a/ble2mqtt/devices/__init__.py b/ble2mqtt/devices/__init__.py
index 818875c..ac778c1 100644
--- a/ble2mqtt/devices/__init__.py
+++ b/ble2mqtt/devices/__init__.py
@@ -9,6 +9,7 @@
 from .presence import Presence  # noqa: F401
 from .qingping_cgdk2 import QingpingTempRHMonitorLite  # noqa: F401
 from .thermostat_ensto import EnstoThermostat  # noqa: F401
+from .voltage_bm2 import VoltageTesterBM2  # noqa: F401
 from .vson_air_wp6003 import VsonWP6003  # noqa: F401
 from .xiaomi_ht import XiaomiHumidityTemperatureV1  # noqa: F401
 from .xiaomi_lywsd03 import XiaomiHumidityTemperatureLYWSD  # noqa: F401
diff --git a/ble2mqtt/devices/voltage_bm2.py b/ble2mqtt/devices/voltage_bm2.py
new file mode 100644
index 0000000..5881d56
--- /dev/null
+++ b/ble2mqtt/devices/voltage_bm2.py
@@ -0,0 +1,56 @@
+from dataclasses import dataclass
+
+from ..devices.base import SENSOR_DOMAIN, Sensor, SubscribeAndSetDataMixin, ConnectionMode
+
+UUID_KEY_READ = "0000fff4-0000-1000-8000-00805f9b34fb"
+KEY = b"\x6c\x65\x61\x67\x65\x6e\x64\xff\xfe\x31\x38\x38\x32\x34\x36\x36"
+
+
+def create_aes():
+    try:
+        from Crypto.Cipher import AES
+    except ImportError:
+        raise ImportError(
+            "Please install pycryptodome to setup BM2 Voltage meter",
+        ) from None
+
+    return AES.new(KEY, AES.MODE_CBC, bytes([0] * 16))
+
+
+@dataclass
+class SensorState:
+    voltage: float
+
+    @classmethod
+    def from_data(cls, decrypted_data: bytes):
+        voltage = (int.from_bytes(
+            decrypted_data[1:1 + 2],
+            byteorder='big',
+        ) >> 4) / 100
+        return cls(voltage=round(voltage, 2))
+
+
+class VoltageTesterBM2(SubscribeAndSetDataMixin, Sensor):
+    NAME = 'voltage_bm2'
+    DATA_CHAR = UUID_KEY_READ
+    MANUFACTURER = 'BM2'
+    SENSOR_CLASS = SensorState
+    REQUIRED_VALUES = ('voltage', )
+    ACTIVE_CONNECTION_MODE = ConnectionMode.ACTIVE_POLL_WITH_DISCONNECT
+
+    @property
+    def entities(self):
+        return {
+            SENSOR_DOMAIN: [
+                {
+                    'name': 'voltage',
+                    'device_class': 'voltage',
+                    'unit_of_measurement': 'V',
+                },
+            ],
+        }
+
+    def process_data(self, data: bytearray):
+        decrypted_data = create_aes().decrypt(data)
+        if decrypted_data[0] == 0xf5:
+            self._state = self.SENSOR_CLASS.from_data(decrypted_data)
diff --git a/requirements.txt b/requirements.txt
index 09f8778..f37e802 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,2 +1,3 @@
 aio-mqtt-mod>=0.2.0
 bleak>=0.12.0
+pycryptodome
diff --git a/setup.py b/setup.py
index 2a30d82..c89613f 100644
--- a/setup.py
+++ b/setup.py
@@ -24,6 +24,9 @@
         'aio-mqtt-mod>=0.3.0',
         'bleak>=0.12.0',
     ],
+    extras_require={
+        'full': ['pycryptodome']
+    },
     classifiers=[
         'Programming Language :: Python :: 3.7',
         'Programming Language :: Python :: 3.8',