feat: add commands to add a new device (#307)

* feat: add commands to add a new device

* chore: mr comments

Author: Lash-L

roborock/web_api.py

@@ -9,7 +9,7 @@
 import time
 
 import aiohttp
-from aiohttp import ContentTypeError
+from aiohttp import ContentTypeError, FormData
 
 from roborock.containers import HomeData, HomeDataRoom, ProductResponse, RRiot, UserData
 from roborock.exceptions import (
@@ -67,22 +67,108 @@ def _get_header_client_id(self):
         md5.update(self._device_identifier.encode())
         return base64.b64encode(md5.digest()).decode()
 
-    def _get_hawk_authentication(self, rriot: RRiot, url: str) -> str:
+    def _process_extra_hawk_values(self, values: dict | None) -> str:
+        if values is None:
+            return ""
+        else:
+            sorted_keys = sorted(values.keys())
+            result = []
+            for key in sorted_keys:
+                value = values.get(key)
+                result.append(f"{key}={value}")
+            return hashlib.md5("&".join(result).encode()).hexdigest()
+
+    def _get_hawk_authentication(
+        self, rriot: RRiot, url: str, formdata: dict | None = None, params: dict | None = None
+    ) -> str:
         timestamp = math.floor(time.time())
         nonce = secrets.token_urlsafe(6)
+        formdata_str = self._process_extra_hawk_values(formdata)
+        params_str = self._process_extra_hawk_values(params)
+
         prestr = ":".join(
             [
                 rriot.u,
                 rriot.s,
                 nonce,
                 str(timestamp),
                 hashlib.md5(url.encode()).hexdigest(),
-                "",
-                "",
+                params_str,
+                formdata_str,
             ]
         )
         mac = base64.b64encode(hmac.new(rriot.h.encode(), prestr.encode(), hashlib.sha256).digest()).decode()
-        return f'Hawk id="{rriot.u}", s="{rriot.s}", ts="{timestamp}", nonce="{nonce}", mac="{mac}"'
+        return f'Hawk id="{rriot.u}",s="{rriot.s}",ts="{timestamp}",nonce="{nonce}",mac="{mac}"'
+
+    async def nc_prepare(self, user_data: UserData, timezone: str) -> dict:
+        """This gets a few critical parameters for adding a device to your account."""
+        if (
+            user_data.rriot is None
+            or user_data.rriot.r is None
+            or user_data.rriot.u is None
+            or user_data.rriot.r.a is None
+        ):
+            raise RoborockException("Your userdata is missing critical attributes.")
+        base_url = user_data.rriot.r.a
+        prepare_request = PreparedRequest(base_url)
+        hid = await self._get_home_id(user_data)
+
+        data = FormData()
+        data.add_field("hid", hid)
+        data.add_field("tzid", timezone)
+
+        prepare_response = await prepare_request.request(
+            "post",
+            "/nc/prepare",
+            headers={
+                "Authorization": self._get_hawk_authentication(
+                    user_data.rriot, "/nc/prepare", {"hid": hid, "tzid": timezone}
+                ),
+            },
+            data=data,
+        )
+
+        if prepare_response is None:
+            raise RoborockException("prepare_response is None")
+        if not prepare_response.get("success"):
+            raise RoborockException(f"{prepare_response.get('msg')} - response code: {prepare_response.get('code')}")
+
+        return prepare_response["result"]
+
+    async def add_device(self, user_data: UserData, s: str, t: str) -> dict:
+        """This will add a new device to your account
+        it is recommended to only use this during a pairing cycle with a device.
+        Please see here: https://github.com/Python-roborock/Roborockmitmproxy/blob/main/handshake_protocol.md
+        """
+        if (
+            user_data.rriot is None
+            or user_data.rriot.r is None
+            or user_data.rriot.u is None
+            or user_data.rriot.r.a is None
+        ):
+            raise RoborockException("Your userdata is missing critical attributes.")
+        base_url = user_data.rriot.r.a
+        add_device_request = PreparedRequest(base_url)
+
+        add_device_response = await add_device_request.request(
+            "GET",
+            "/user/devices/newadd",
+            headers={
+                "Authorization": self._get_hawk_authentication(
+                    user_data.rriot, "/user/devices/newadd", params={"s": s, "t": t}
+                ),
+            },
+            params={"s": s, "t": t},
+        )
+
+        if add_device_response is None:
+            raise RoborockException("add_device is None")
+        if not add_device_response.get("success"):
+            raise RoborockException(
+                f"{add_device_response.get('msg')} - response code: {add_device_response.get('code')}"
+            )
+
+        return add_device_response["result"]
 
     async def request_code(self) -> None:
         base_url = await self._get_base_url()

tests/conftest.py

@@ -200,6 +200,57 @@ def mock_rest() -> aioresponses:
             status=200,
             payload={"api": None, "code": 200, "result": HOME_DATA_RAW, "status": "ok", "success": True},
         )
+        mocked.post(
+            re.compile(r"https://api-.*\.roborock\.com/nc/prepare"),
+            status=200,
+            payload={
+                "api": None,
+                "result": {"r": "US", "s": "ffffff", "t": "eOf6d2BBBB"},
+                "status": "ok",
+                "success": True,
+            },
+        )
+
+        mocked.get(
+            re.compile(r"https://api-.*\.roborock\.com/user/devices/newadd/*"),
+            status=200,
+            payload={
+                "api": "获取新增设备信息",
+                "result": {
+                    "activeTime": 1737724598,
+                    "attribute": None,
+                    "cid": None,
+                    "createTime": 0,
+                    "deviceStatus": None,
+                    "duid": "rand_duid",
+                    "extra": "{}",
+                    "f": False,
+                    "featureSet": "0",
+                    "fv": "02.16.12",
+                    "iconUrl": "",
+                    "lat": None,
+                    "localKey": "random_lk",
+                    "lon": None,
+                    "name": "S7",
+                    "newFeatureSet": "0000000000002000",
+                    "online": True,
+                    "productId": "rand_prod_id",
+                    "pv": "1.0",
+                    "roomId": None,
+                    "runtimeEnv": None,
+                    "setting": None,
+                    "share": False,
+                    "shareTime": None,
+                    "silentOtaSwitch": False,
+                    "sn": "Rand_sn",
+                    "timeZoneId": "America/New_York",
+                    "tuyaMigrated": False,
+                    "tuyaUuid": None,
+                },
+                "status": "ok",
+                "success": True,
+            },
+        )
         yield mocked
 
 

tests/test_web_api.py

@@ -26,3 +26,13 @@ async def test_get_home_data_v2():
     ud = await api.code_login(4123)
     hd = await api.get_home_data_v2(ud)
     assert hd == HomeData.from_dict(HOME_DATA_RAW)
+
+
+async def test_nc_prepare():
+    """Test adding a device and that nothing breaks"""
+    api = RoborockApiClient(username="[email protected]")
+    await api.request_code()
+    ud = await api.code_login(4123)
+    prepare = await api.nc_prepare(ud, "America/New_York")
+    new_device = await api.add_device(ud, prepare["s"], prepare["t"])
+    assert new_device["duid"] == "rand_duid"