Model whitelist

6 files changed, 38 insertions, 43 deletions

diff --git a/custom_components/skykettle/__init__.py b/custom_components/skykettle/__init__.py
index 0f937d4..4dc2092 100644
--- a/custom_components/skykettle/__init__.py
+++ b/custom_components/skykettle/__init__.py
@@ -28,12 +28,20 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry):
     if DOMAIN not in hass.data: hass.data[DOMAIN] = {}
     if entry.entry_id not in hass.data: hass.data[DOMAIN][entry.entry_id] = {}
 
+    # Backward compatibility
+    model = entry.data.get(CONF_FRIENDLY_NAME, None)
+    if model != None and not model.endswith("S"):
+        config = dict(entry.data.items())
+        config[CONF_FRIENDLY_NAME] = model + "S"
+        hass.config_entries.async_update_entry(entry, data=config)
+
     kettle = KettleConnection(
         mac=entry.data[CONF_MAC],
         key=entry.data[CONF_PASSWORD],
         persistent=entry.data[CONF_PERSISTENT_CONNECTION],
         adapter=entry.data.get(CONF_DEVICE, None),
-        hass=hass
+        hass=hass,
+        model=entry.data.get(CONF_FRIENDLY_NAME, None)
     )
     hass.data[DOMAIN][entry.entry_id][DATA_CONNECTION] = kettle
 
diff --git a/custom_components/skykettle/config_flow.py b/custom_components/skykettle/config_flow.py
index 9721a47..419037e 100644
--- a/custom_components/skykettle/config_flow.py
+++ b/custom_components/skykettle/config_flow.py
@@ -14,6 +14,7 @@ from .const import *
 from .ble_scan import ble_scan
 from .ble_scan import ble_get_adapters
 from .kettle_connection import KettleConnection
+from .skykettle import SkyKettle
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -112,26 +113,23 @@ class SkyKettleConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         """Handle the scan step."""
         errors = {}
         if user_input is not None:
-            if user_input[CONF_MAC] == "...": return await self.async_step_manual_mac()
             spl = user_input[CONF_MAC].split(' ', maxsplit=1)
             mac = spl[0]
-            name = spl[1][1:-2] if len(spl) >= 2 else None
-            if not await self.init_mac(mac): return self.async_abort(reason='already_configured')
+            name = spl[1][1:-1] if len(spl) >= 2 else None
+            if not await self.init_mac(mac):
+                # This kettle already configured
+                return self.async_abort(reason='already_configured')
             if name: self.config[CONF_FRIENDLY_NAME] = name
-            # Continue to options
+            # Continue to connect step
             return await self.async_step_connect()
 
         try:
             macs = await ble_scan(self.config.get(CONF_DEVICE, None), scan_time=BLE_SCAN_TIME)
             _LOGGER.debug(f"Scan result: {macs}")
-            macs_filtered = [mac for mac in macs if mac.name and mac.name.startswith("RK-")]
-            if len(macs_filtered) > 0:
-                macs = macs_filtered
-                _LOGGER.debug(f"Filtered scan result: {macs}")
-            if len(macs) == 0:
-                errors["base"] = "no_scan"
-            mac_list = [f"{r.mac} ({r.name})" for r in macs]
-            mac_list = mac_list + ["..."] # Manual
+            macs_filtered = [mac for mac in macs if mac.name and mac.name in SkyKettle.SUPPORTED_DEVICES]
+            if len(macs_filtered) == 0:
+                return self.async_abort(reason='kettle_not_found')
+            mac_list = [f"{r.mac} ({r.name})" for r in macs_filtered]
             schema = vol.Schema(
             {
                 vol.Required(CONF_MAC): vol.In(mac_list)
@@ -151,29 +149,6 @@ class SkyKettleConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
             data_schema=schema
         )
 
-    async def async_step_manual_mac(self, user_input=None):
-        """Handle the manual_mac step."""
-        errors = {}
-        if user_input is not None:
-            mac = user_input[CONF_MAC]
-            if re.match(REGEX_MAC, mac):
-                if not await self.init_mac(mac): return self.async_abort(reason='already_configured')
-            else:
-                errors[CONF_MAC] = "invalid_mac"
-            if not errors:
-                # Continue to options
-                return await self.async_step_connect()
-
-        MAC_SCHEMA = vol.Schema(
-        {
-            vol.Required(CONF_MAC): cv.string,
-        })
-        return self.async_show_form(
-            step_id="manual_mac",
-            errors=errors,
-            data_schema=MAC_SCHEMA
-        )
-
     async def async_step_connect(self, user_input=None):
         """Handle the connect step."""
         errors = {}
diff --git a/custom_components/skykettle/kettle_connection.py b/custom_components/skykettle/kettle_connection.py
index 74b88e3..d3ec9a5 100644
--- a/custom_components/skykettle/kettle_connection.py
+++ b/custom_components/skykettle/kettle_connection.py
@@ -22,8 +22,8 @@ class KettleConnection(SkyKettle):
     STATS_INTERVAL = 15
     TARGET_TTL = 30
 
-    def __init__(self, mac, key, persistent=True, adapter=None, hass=None):
-        super().__init__()
+    def __init__(self, mac, key, persistent=True, adapter=None, hass=None, model=None):
+        super().__init__(model)
         self._child = None
         self._mac = mac
         self._key = key
@@ -164,8 +164,8 @@ class KettleConnection(SkyKettle):
                 _LOGGER.warning(f"Auth failed. You need to enable pairing mode on the kettle.")
                 raise AuthError("Auth failed")
             _LOGGER.debug("Auth ok")
-            await self.sync_time()
             self._sw_version = await self.get_version()
+            await self.sync_time()
 
     async def _disconnect_if_need(self):
         if not self.persistent and self.target_mode != SkyKettle.MODE_GAME:
diff --git a/custom_components/skykettle/skykettle.py b/custom_components/skykettle/skykettle.py
index 2dbcaf6..2f48a99 100644
--- a/custom_components/skykettle/skykettle.py
+++ b/custom_components/skykettle/skykettle.py
@@ -10,6 +10,9 @@ _LOGGER = logging.getLogger(__name__)
 
 
 class SkyKettle():
+    # Source: https://github.com/mavrikkk/ha_kettler/blob/master/custom_components/ready4sky/r4sconst.py
+    SUPPORTED_DEVICES = {'RK-M170S':0, 'RK-M171S':0, 'RK-M173S':0, 'RK-G200S':1, 'RK-G201S':1, 'RK-G202S':1, 'RK-G210S':1, 'RK-G211S':1, 'RK-G212S':1, 'RK-G240S':1, 'RK-M216S':2, 'RK-M216S-E':2}
+
     MODE_BOIL = 0x00
     MODE_HEAT = 0x01
     MODE_BOIL_HEAT = 0x02
@@ -75,6 +78,17 @@ class SkyKettle():
     Stats = namedtuple("Stats", ["ontime", "energy_wh", "heater_on_count", "user_on_count"])
     FreshWaterInfo = namedtuple("FreshWaterInfo", ["is_on", "unknown1", "water_freshness_hours"])
 
+
+    def __init__(self, model):
+        _LOGGER.info(f"Kettle model: {model}")
+        self.model = model
+        if model in SkyKettle.SUPPORTED_DEVICES:
+            self.model_code = SkyKettle.SUPPORTED_DEVICES[model]
+            _LOGGER.debug(f"Kettle model code: {self.model_code}")
+        else:
+            self.model_code = 1
+            _LOGGER.warn(f"Unknown kettle model")
+
     @abstractmethod
     async def command(self, command, params=[]):
         pass
diff --git a/custom_components/skykettle/translations/en.json b/custom_components/skykettle/translations/en.json
index 623d928..fb5d68e 100644
--- a/custom_components/skykettle/translations/en.json
+++ b/custom_components/skykettle/translations/en.json
@@ -2,6 +2,7 @@
     "title": "SkyKettle",
     "config": {
         "abort": {
+            "kettle_not_found": "Can't find any supported kettle. Please make sure that your kettle is in the list of supported devices and it's on the stand.",
             "already_configured": "Integration for this kettle is already configured.",
             "linux_not_found": "Sorry, Linux-based system required.",
             "hcitool_not_found": "hcitool binary not found.",
@@ -11,8 +12,6 @@
             "unknown": "Unknown error, please check logs."
         },
         "error": {
-            "no_scan": "I can't find any BLE device for some reason, you can enter MAC address manually on the next step.",
-            "invalid_mac": "MAC address must be in the form \"AA:BB:CC:DD:EE:FF\".",
             "cant_connect": "Can't connect to the kettle. Please make sure that it's on the stand, hatttool is granted with privileges (check instrucrions) and official application is not connected to the kettle.",
             "cant_auth": "Connected to the kettle but pairing failed. Please try again."
         },
diff --git a/custom_components/skykettle/translations/ru.json b/custom_components/skykettle/translations/ru.json
index 2d24b1c..7f87fd7 100644
--- a/custom_components/skykettle/translations/ru.json
+++ b/custom_components/skykettle/translations/ru.json
@@ -2,6 +2,7 @@
     "title": "SkyKettle",
     "config": {
         "abort": {
+            "kettle_not_found": "Не могу найти чайник. Пожалуйста, убедитесь, что модель вашего чайника есть в списке поддерживаемых, и что он стоит на подставке.",
             "already_configured": "Для этого чайника интеграция уже настроена.",
             "linux_not_found": "Извините, эта интеграция может работать только под Linux.",
             "hcitool_not_found": "Утилита \"hcitool\" не найдена.",
@@ -11,8 +12,6 @@
             "unknown": "Неизвестная ошибка, смотрите логи."
         },
         "error": {
-            "no_scan": "Почему-то не могу найти ни одного устройства, попробуйте на следующем шаге ввести MAC-адрес вручную.",
-            "invalid_mac": "MAC-адрес должен быть в формате \"AA:BB:CC:DD:EE:FF\".",
             "cant_connect": "Не могу подключиться к чайнику. Пожалуйста, убедитесь, что он стоит на подставке, включен в розетку, утилите \"hatttool\" предоставлены права (смотрите инструкцию), и официальное приложение сейчас не подключено к чайнику.",
             "cant_auth": "Подключились к чайнику, но сопряжение не удалось. Попробуйте снова."
         },