1from __future__ import annotations
  2
  3import base64
  4import hashlib
  5import hmac
  6import logging
  7import math
  8import secrets
  9import string
 10import time
 11from dataclasses import dataclass
 12
 13import aiohttp
 14from aiohttp import ContentTypeError, FormData
 15from pyrate_limiter import BucketFullException, Duration, Limiter, Rate
 16
 17from roborock import HomeDataSchedule
 18from roborock.data import HomeData, HomeDataRoom, HomeDataScene, ProductResponse, RRiot, UserData
 19from roborock.exceptions import (
 20    RoborockAccountDoesNotExist,
 21    RoborockException,
 22    RoborockInvalidCode,
 23    RoborockInvalidCredentials,
 24    RoborockInvalidEmail,
 25    RoborockInvalidUserAgreement,
 26    RoborockMissingParameters,
 27    RoborockNoResponseFromBaseURL,
 28    RoborockNoUserAgreement,
 29    RoborockRateLimit,
 30    RoborockTooFrequentCodeRequests,
 31)
 32
 33_LOGGER = logging.getLogger(__name__)
 34BASE_URLS = [
 35    "https://usiot.roborock.com",
 36    "https://euiot.roborock.com",
 37    "https://cniot.roborock.com",
 38    "https://ruiot.roborock.com",
 39]
 40
 41
 42@dataclass
 43class IotLoginInfo:
 44    """Information about the login to the iot server."""
 45
 46    base_url: str
 47    country_code: str
 48    country: str
 49
 50
 51class RoborockApiClient:
 52    _LOGIN_RATES = [
 53        Rate(1, Duration.SECOND),
 54        Rate(3, Duration.MINUTE),
 55        Rate(10, Duration.HOUR),
 56        Rate(20, Duration.DAY),
 57    ]
 58    _HOME_DATA_RATES = [
 59        Rate(1, Duration.SECOND),
 60        Rate(5, Duration.MINUTE),
 61        Rate(15, Duration.HOUR),
 62        Rate(40, Duration.DAY),
 63    ]
 64
 65    _login_limiter = Limiter(_LOGIN_RATES)
 66    _home_data_limiter = Limiter(_HOME_DATA_RATES)
 67
 68    def __init__(
 69        self, username: str, base_url: str | None = None, session: aiohttp.ClientSession | None = None
 70    ) -> None:
 71        """Sample API Client."""
 72        self._username = username
 73        self._base_url = base_url
 74        self._device_identifier = secrets.token_urlsafe(16)
 75        self.session = session
 76        self._iot_login_info: IotLoginInfo | None = None
 77
 78    async def _get_iot_login_info(self) -> IotLoginInfo:
 79        if self._iot_login_info is None:
 80            valid_urls = BASE_URLS if self._base_url is None else [self._base_url]
 81            for iot_url in valid_urls:
 82                url_request = PreparedRequest(iot_url, self.session)
 83                response = await url_request.request(
 84                    "post",
 85                    "/api/v1/getUrlByEmail",
 86                    params={"email": self._username, "needtwostepauth": "false"},
 87                )
 88                if response is None:
 89                    continue
 90                response_code = response.get("code")
 91                if response_code != 200:
 92                    if response_code == 2003:
 93                        raise RoborockInvalidEmail("Your email was incorrectly formatted.")
 94                    elif response_code == 1001:
 95                        raise RoborockMissingParameters(
 96                            "You are missing parameters for this request, are you sure you entered your username?"
 97                        )
 98                    else:
 99                        raise RoborockException(f"{response.get('msg')} - response code: {response_code}")
100                country_code = response["data"]["countrycode"]
101                country = response["data"]["country"]
102                if country_code is not None or country is not None:
103                    self._iot_login_info = IotLoginInfo(
104                        base_url=response["data"]["url"],
105                        country=country,
106                        country_code=country_code,
107                    )
108                    _LOGGER.debug("Country determined to be %s and code is %s", country, country_code)
109                    return self._iot_login_info
110            raise RoborockNoResponseFromBaseURL(
111                "No account was found for any base url we tried. Either your email is incorrect or we do not have a"
112                " record of the roborock server your device is on."
113            )
114        return self._iot_login_info
115
116    @property
117    async def base_url(self):
118        if self._base_url is not None:
119            return self._base_url
120        return (await self._get_iot_login_info()).base_url
121
122    @property
123    async def country(self):
124        return (await self._get_iot_login_info()).country
125
126    @property
127    async def country_code(self):
128        return (await self._get_iot_login_info()).country_code
129
130    def _get_header_client_id(self):
131        md5 = hashlib.md5()
132        md5.update(self._username.encode())
133        md5.update(self._device_identifier.encode())
134        return base64.b64encode(md5.digest()).decode()
135
136    async def nc_prepare(self, user_data: UserData, timezone: str) -> dict:
137        """This gets a few critical parameters for adding a device to your account."""
138        if (
139            user_data.rriot is None
140            or user_data.rriot.r is None
141            or user_data.rriot.u is None
142            or user_data.rriot.r.a is None
143        ):
144            raise RoborockException("Your userdata is missing critical attributes.")
145        base_url = user_data.rriot.r.a
146        prepare_request = PreparedRequest(base_url, self.session)
147        hid = await self._get_home_id(user_data)
148
149        data = FormData()
150        data.add_field("hid", hid)
151        data.add_field("tzid", timezone)
152
153        prepare_response = await prepare_request.request(
154            "post",
155            "/nc/prepare",
156            headers={
157                "Authorization": _get_hawk_authentication(
158                    user_data.rriot, "/nc/prepare", {"hid": hid, "tzid": timezone}
159                ),
160            },
161            data=data,
162        )
163
164        if prepare_response is None:
165            raise RoborockException("prepare_response is None")
166        if not prepare_response.get("success"):
167            raise RoborockException(f"{prepare_response.get('msg')} - response code: {prepare_response.get('code')}")
168
169        return prepare_response["result"]
170
171    async def add_device(self, user_data: UserData, s: str, t: str) -> dict:
172        """This will add a new device to your account
173        it is recommended to only use this during a pairing cycle with a device.
174        Please see here: https://github.com/Python-roborock/Roborockmitmproxy/blob/main/handshake_protocol.md
175        """
176        if (
177            user_data.rriot is None
178            or user_data.rriot.r is None
179            or user_data.rriot.u is None
180            or user_data.rriot.r.a is None
181        ):
182            raise RoborockException("Your userdata is missing critical attributes.")
183        base_url = user_data.rriot.r.a
184        add_device_request = PreparedRequest(base_url, self.session)
185
186        add_device_response = await add_device_request.request(
187            "GET",
188            "/user/devices/newadd",
189            headers={
190                "Authorization": _get_hawk_authentication(
191                    user_data.rriot, "/user/devices/newadd", params={"s": s, "t": t}
192                ),
193            },
194            params={"s": s, "t": t},
195        )
196
197        if add_device_response is None:
198            raise RoborockException("add_device is None")
199        if not add_device_response.get("success"):
200            raise RoborockException(
201                f"{add_device_response.get('msg')} - response code: {add_device_response.get('code')}"
202            )
203
204        return add_device_response["result"]
205
206    async def request_code(self) -> None:
207        try:
208            self._login_limiter.try_acquire("login")
209        except BucketFullException as ex:
210            _LOGGER.info(ex.meta_info)
211            raise RoborockRateLimit("Reached maximum requests for login. Please try again later.") from ex
212        base_url = await self.base_url
213        header_clientid = self._get_header_client_id()
214        code_request = PreparedRequest(base_url, self.session, {"header_clientid": header_clientid})
215
216        code_response = await code_request.request(
217            "post",
218            "/api/v1/sendEmailCode",
219            params={
220                "username": self._username,
221                "type": "auth",
222            },
223        )
224        if code_response is None:
225            raise RoborockException("Failed to get a response from send email code")
226        response_code = code_response.get("code")
227        if response_code != 200:
228            _LOGGER.info("Request code failed for %s with the following context: %s", self._username, code_response)
229            if response_code == 2008:
230                raise RoborockAccountDoesNotExist("Account does not exist - check your login and try again.")
231            elif response_code == 9002:
232                raise RoborockTooFrequentCodeRequests("You have attempted to request too many codes. Try again later")
233            else:
234                raise RoborockException(f"{code_response.get('msg')} - response code: {code_response.get('code')}")
235
236    async def request_code_v4(self) -> None:
237        """Request a code using the v4 endpoint."""
238        if await self.country_code is None or await self.country is None:
239            _LOGGER.info("No country code or country found, trying old version of request code.")
240            return await self.request_code()
241        try:
242            self._login_limiter.try_acquire("login")
243        except BucketFullException as ex:
244            _LOGGER.info(ex.meta_info)
245            raise RoborockRateLimit("Reached maximum requests for login. Please try again later.") from ex
246        base_url = await self.base_url
247        header_clientid = self._get_header_client_id()
248        code_request = PreparedRequest(
249            base_url,
250            self.session,
251            {
252                "header_clientid": header_clientid,
253                "Content-Type": "application/x-www-form-urlencoded",
254                "header_clientlang": "en",
255            },
256        )
257
258        code_response = await code_request.request(
259            "post",
260            "/api/v4/email/code/send",
261            data={"email": self._username, "type": "login", "platform": ""},
262        )
263        if code_response is None:
264            raise RoborockException("Failed to get a response from send email code")
265        response_code = code_response.get("code")
266        if response_code != 200:
267            _LOGGER.info("Request code failed for %s with the following context: %s", self._username, code_response)
268            if response_code == 2008:
269                raise RoborockAccountDoesNotExist("Account does not exist - check your login and try again.")
270            elif response_code == 9002:
271                raise RoborockTooFrequentCodeRequests("You have attempted to request too many codes. Try again later")
272            else:
273                raise RoborockException(f"{code_response.get('msg')} - response code: {code_response.get('code')}")
274
275    async def _sign_key_v3(self, s: str) -> str:
276        """Sign a randomly generated string."""
277        base_url = await self.base_url
278        header_clientid = self._get_header_client_id()
279        code_request = PreparedRequest(base_url, self.session, {"header_clientid": header_clientid})
280
281        code_response = await code_request.request(
282            "post",
283            "/api/v3/key/sign",
284            params={"s": s},
285        )
286
287        if not code_response or "data" not in code_response or "k" not in code_response["data"]:
288            raise RoborockException("Failed to get a response from sign key")
289        response_code = code_response.get("code")
290
291        if response_code != 200:
292            _LOGGER.info("Request code failed for %s with the following context: %s", self._username, code_response)
293            raise RoborockException(f"{code_response.get('msg')} - response code: {code_response.get('code')}")
294
295        return code_response["data"]["k"]
296
297    async def code_login_v4(
298        self, code: int | str, country: str | None = None, country_code: int | None = None
299    ) -> UserData:
300        """
301        Login via code authentication.
302        :param code: The code from the email.
303        :param country: The two-character representation of the country, i.e. "US"
304        :param country_code: the country phone number code i.e. 1 for US.
305        """
306        base_url = await self.base_url
307        if country is None:
308            country = await self.country
309        if country_code is None:
310            country_code = await self.country_code
311        if country_code is None or country is None:
312            _LOGGER.info("No country code or country found, trying old version of code login.")
313            return await self.code_login(code)
314        header_clientid = self._get_header_client_id()
315        x_mercy_ks = "".join(secrets.choice(string.ascii_letters + string.digits) for _ in range(16))
316        x_mercy_k = await self._sign_key_v3(x_mercy_ks)
317        login_request = PreparedRequest(
318            base_url,
319            self.session,
320            {
321                "header_clientid": header_clientid,
322                "x-mercy-ks": x_mercy_ks,
323                "x-mercy-k": x_mercy_k,
324                "Content-Type": "application/x-www-form-urlencoded",
325                "header_clientlang": "en",
326                "header_appversion": "4.54.02",
327                "header_phonesystem": "iOS",
328                "header_phonemodel": "iPhone16,1",
329            },
330        )
331        login_response = await login_request.request(
332            "post",
333            "/api/v4/auth/email/login/code",
334            data={
335                "country": country,
336                "countryCode": country_code,
337                "email": self._username,
338                "code": code,
339                # Major and minor version are the user agreement version, we will need to see if this needs to be
340                # dynamic https://usiot.roborock.com/api/v3/app/agreement/latest?country=US
341                "majorVersion": 14,
342                "minorVersion": 0,
343            },
344        )
345        if login_response is None:
346            raise RoborockException("Login request response is None")
347        response_code = login_response.get("code")
348        if response_code != 200:
349            _LOGGER.info("Login failed for %s with the following context: %s", self._username, login_response)
350            if response_code == 2018:
351                raise RoborockInvalidCode("Invalid code - check your code and try again.")
352            if response_code == 3009:
353                raise RoborockNoUserAgreement("You must accept the user agreement in the Roborock app to continue.")
354            if response_code == 3006:
355                raise RoborockInvalidUserAgreement(
356                    "User agreement must be accepted again - or you are attempting to use the Mi Home app account."
357                )
358            raise RoborockException(f"{login_response.get('msg')} - response code: {response_code}")
359        user_data = login_response.get("data")
360        if not isinstance(user_data, dict):
361            raise RoborockException("Got unexpected data type for user_data")
362        return UserData.from_dict(user_data)
363
364    async def pass_login(self, password: str) -> UserData:
365        try:
366            self._login_limiter.try_acquire("login")
367        except BucketFullException as ex:
368            _LOGGER.info(ex.meta_info)
369            raise RoborockRateLimit("Reached maximum requests for login. Please try again later.") from ex
370        base_url = await self.base_url
371        header_clientid = self._get_header_client_id()
372
373        login_request = PreparedRequest(base_url, self.session, {"header_clientid": header_clientid})
374        login_response = await login_request.request(
375            "post",
376            "/api/v1/login",
377            params={
378                "username": self._username,
379                "password": password,
380                "needtwostepauth": "false",
381            },
382        )
383        if login_response is None:
384            raise RoborockException("Login response is none")
385        if login_response.get("code") != 200:
386            _LOGGER.info("Login failed for %s with the following context: %s", self._username, login_response)
387            raise RoborockException(f"{login_response.get('msg')} - response code: {login_response.get('code')}")
388        user_data = login_response.get("data")
389        if not isinstance(user_data, dict):
390            raise RoborockException("Got unexpected data type for user_data")
391        return UserData.from_dict(user_data)
392
393    async def pass_login_v3(self, password: str) -> UserData:
394        """Seemingly it follows the format below, but password is encrypted in some manner.
395        # login_response = await login_request.request(
396        #     "post",
397        #     "/api/v3/auth/email/login",
398        #     params={
399        #         "email": self._username,
400        #         "password": password,
401        #         "twoStep": 1,
402        #         "version": 0
403        #     },
404        # )
405        """
406        raise NotImplementedError("Pass_login_v3 has not yet been implemented")
407
408    async def code_login(self, code: int | str) -> UserData:
409        base_url = await self.base_url
410        header_clientid = self._get_header_client_id()
411
412        login_request = PreparedRequest(base_url, self.session, {"header_clientid": header_clientid})
413        login_response = await login_request.request(
414            "post",
415            "/api/v1/loginWithCode",
416            params={
417                "username": self._username,
418                "verifycode": code,
419                "verifycodetype": "AUTH_EMAIL_CODE",
420            },
421        )
422        if login_response is None:
423            raise RoborockException("Login request response is None")
424        response_code = login_response.get("code")
425        if response_code != 200:
426            _LOGGER.info("Login failed for %s with the following context: %s", self._username, login_response)
427            if response_code == 2018:
428                raise RoborockInvalidCode("Invalid code - check your code and try again.")
429            if response_code == 3009:
430                raise RoborockNoUserAgreement("You must accept the user agreement in the Roborock app to continue.")
431            if response_code == 3006:
432                raise RoborockInvalidUserAgreement(
433                    "User agreement must be accepted again - or you are attempting to use the Mi Home app account."
434                )
435            raise RoborockException(f"{login_response.get('msg')} - response code: {response_code}")
436        user_data = login_response.get("data")
437        if not isinstance(user_data, dict):
438            raise RoborockException("Got unexpected data type for user_data")
439        return UserData.from_dict(user_data)
440
441    async def _get_home_id(self, user_data: UserData):
442        base_url = await self.base_url
443        header_clientid = self._get_header_client_id()
444        home_id_request = PreparedRequest(base_url, self.session, {"header_clientid": header_clientid})
445        home_id_response = await home_id_request.request(
446            "get",
447            "/api/v1/getHomeDetail",
448            headers={"Authorization": user_data.token},
449        )
450        if home_id_response is None:
451            raise RoborockException("home_id_response is None")
452        if home_id_response.get("code") != 200:
453            _LOGGER.info("Get Home Id failed with the following context: %s", home_id_response)
454            if home_id_response.get("code") == 2010:
455                raise RoborockInvalidCredentials(
456                    f"Invalid credentials ({home_id_response.get('msg')}) - check your login and try again."
457                )
458            raise RoborockException(f"{home_id_response.get('msg')} - response code: {home_id_response.get('code')}")
459
460        return home_id_response["data"]["rrHomeId"]
461
462    async def get_home_data(self, user_data: UserData) -> HomeData:
463        try:
464            self._home_data_limiter.try_acquire("home_data")
465        except BucketFullException as ex:
466            _LOGGER.info(ex.meta_info)
467            raise RoborockRateLimit("Reached maximum requests for home data. Please try again later.") from ex
468        rriot = user_data.rriot
469        if rriot is None:
470            raise RoborockException("rriot is none")
471        home_id = await self._get_home_id(user_data)
472        if rriot.r.a is None:
473            raise RoborockException("Missing field 'a' in rriot reference")
474        home_request = PreparedRequest(
475            rriot.r.a,
476            self.session,
477            {
478                "Authorization": _get_hawk_authentication(rriot, f"/user/homes/{str(home_id)}"),
479            },
480        )
481        home_response = await home_request.request("get", "/user/homes/" + str(home_id))
482        if not home_response.get("success"):
483            raise RoborockException(home_response)
484        home_data = home_response.get("result")
485        if isinstance(home_data, dict):
486            return HomeData.from_dict(home_data)
487        else:
488            raise RoborockException("home_response result was an unexpected type")
489
490    async def get_home_data_v2(self, user_data: UserData) -> HomeData:
491        """This is the same as get_home_data, but uses a different endpoint and includes non-robotic vacuums."""
492        try:
493            self._home_data_limiter.try_acquire("home_data")
494        except BucketFullException as ex:
495            _LOGGER.info(ex.meta_info)
496            raise RoborockRateLimit("Reached maximum requests for home data. Please try again later.") from ex
497        rriot = user_data.rriot
498        if rriot is None:
499            raise RoborockException("rriot is none")
500        home_id = await self._get_home_id(user_data)
501        if rriot.r.a is None:
502            raise RoborockException("Missing field 'a' in rriot reference")
503        home_request = PreparedRequest(
504            rriot.r.a,
505            self.session,
506            {
507                "Authorization": _get_hawk_authentication(rriot, "/v2/user/homes/" + str(home_id)),
508            },
509        )
510        home_response = await home_request.request("get", "/v2/user/homes/" + str(home_id))
511        if not home_response.get("success"):
512            raise RoborockException(home_response)
513        home_data = home_response.get("result")
514        if isinstance(home_data, dict):
515            return HomeData.from_dict(home_data)
516        else:
517            raise RoborockException("home_response result was an unexpected type")
518
519    async def get_home_data_v3(self, user_data: UserData) -> HomeData:
520        """This is the same as get_home_data, but uses a different endpoint and includes non-robotic vacuums."""
521        try:
522            self._home_data_limiter.try_acquire("home_data")
523        except BucketFullException as ex:
524            _LOGGER.info(ex.meta_info)
525            raise RoborockRateLimit("Reached maximum requests for home data. Please try again later.") from ex
526        rriot = user_data.rriot
527        home_id = await self._get_home_id(user_data)
528        if rriot.r.a is None:
529            raise RoborockException("Missing field 'a' in rriot reference")
530        home_request = PreparedRequest(
531            rriot.r.a,
532            self.session,
533            {
534                "Authorization": _get_hawk_authentication(rriot, "/v3/user/homes/" + str(home_id)),
535            },
536        )
537        home_response = await home_request.request("get", "/v3/user/homes/" + str(home_id))
538        if not home_response.get("success"):
539            raise RoborockException(home_response)
540        home_data = home_response.get("result")
541        if isinstance(home_data, dict):
542            return HomeData.from_dict(home_data)
543        raise RoborockException(f"home_response result was an unexpected type: {home_data}")
544
545    async def get_rooms(self, user_data: UserData, home_id: int | None = None) -> list[HomeDataRoom]:
546        rriot = user_data.rriot
547        if rriot is None:
548            raise RoborockException("rriot is none")
549        if home_id is None:
550            home_id = await self._get_home_id(user_data)
551        if rriot.r.a is None:
552            raise RoborockException("Missing field 'a' in rriot reference")
553        room_request = PreparedRequest(
554            rriot.r.a,
555            self.session,
556            {
557                "Authorization": _get_hawk_authentication(rriot, "/v2/user/homes/" + str(home_id)),
558            },
559        )
560        room_response = await room_request.request("get", f"/user/homes/{str(home_id)}/rooms" + str(home_id))
561        if not room_response.get("success"):
562            raise RoborockException(room_response)
563        rooms = room_response.get("result")
564        if isinstance(rooms, list):
565            output_list = []
566            for room in rooms:
567                output_list.append(HomeDataRoom.from_dict(room))
568            return output_list
569        else:
570            raise RoborockException("home_response result was an unexpected type")
571
572    async def get_scenes(self, user_data: UserData, device_id: str) -> list[HomeDataScene]:
573        rriot = user_data.rriot
574        if rriot is None:
575            raise RoborockException("rriot is none")
576        if rriot.r.a is None:
577            raise RoborockException("Missing field 'a' in rriot reference")
578        scenes_request = PreparedRequest(
579            rriot.r.a,
580            self.session,
581            {
582                "Authorization": _get_hawk_authentication(rriot, f"/user/scene/device/{str(device_id)}"),
583            },
584        )
585        scenes_response = await scenes_request.request("get", f"/user/scene/device/{str(device_id)}")
586        if not scenes_response.get("success"):
587            raise RoborockException(scenes_response)
588        scenes = scenes_response.get("result")
589        if isinstance(scenes, list):
590            return [HomeDataScene.from_dict(scene) for scene in scenes]
591        else:
592            raise RoborockException("scene_response result was an unexpected type")
593
594    async def execute_scene(self, user_data: UserData, scene_id: int) -> None:
595        rriot = user_data.rriot
596        if rriot is None:
597            raise RoborockException("rriot is none")
598        if rriot.r.a is None:
599            raise RoborockException("Missing field 'a' in rriot reference")
600        execute_scene_request = PreparedRequest(
601            rriot.r.a,
602            self.session,
603            {
604                "Authorization": _get_hawk_authentication(rriot, f"/user/scene/{str(scene_id)}/execute"),
605            },
606        )
607        execute_scene_response = await execute_scene_request.request("POST", f"/user/scene/{str(scene_id)}/execute")
608        if not execute_scene_response.get("success"):
609            raise RoborockException(execute_scene_response)
610
611    async def get_schedules(self, user_data: UserData, device_id: str) -> list[HomeDataSchedule]:
612        rriot = user_data.rriot
613        if rriot is None:
614            raise RoborockException("rriot is none")
615        if rriot.r.a is None:
616            raise RoborockException("Missing field 'a' in rriot reference")
617        schedules_request = PreparedRequest(
618            rriot.r.a,
619            self.session,
620            {
621                "Authorization": _get_hawk_authentication(rriot, f"/user/devices/{device_id}/jobs"),
622            },
623        )
624        schedules_response = await schedules_request.request("get", f"/user/devices/{str(device_id)}/jobs")
625        if not schedules_response.get("success"):
626            raise RoborockException(schedules_response)
627        schedules = schedules_response.get("result")
628        if isinstance(schedules, list):
629            return [HomeDataSchedule.from_dict(schedule) for schedule in schedules]
630        else:
631            raise RoborockException(f"schedule_response result was an unexpected type: {schedules}")
632
633    async def get_products(self, user_data: UserData) -> ProductResponse:
634        """Gets all products and their schemas, good for determining status codes and model numbers."""
635        base_url = await self.base_url
636        header_clientid = self._get_header_client_id()
637        product_request = PreparedRequest(base_url, self.session, {"header_clientid": header_clientid})
638        product_response = await product_request.request(
639            "get",
640            "/api/v4/product",
641            headers={"Authorization": user_data.token},
642        )
643        if product_response is None:
644            raise RoborockException("home_id_response is None")
645        if product_response.get("code") != 200:
646            raise RoborockException(f"{product_response.get('msg')} - response code: {product_response.get('code')}")
647        result = product_response.get("data")
648        if isinstance(result, dict):
649            return ProductResponse.from_dict(result)
650        raise RoborockException("product result was an unexpected type")
651
652    async def download_code(self, user_data: UserData, product_id: int):
653        base_url = await self.base_url
654        header_clientid = self._get_header_client_id()
655        product_request = PreparedRequest(base_url, self.session, {"header_clientid": header_clientid})
656        request = {"apilevel": 99999, "productids": [product_id], "type": 2}
657        response = await product_request.request(
658            "post",
659            "/api/v1/appplugin",
660            json=request,
661            headers={"Authorization": user_data.token, "Content-Type": "application/json"},
662        )
663        return response["data"][0]["url"]
664
665    async def download_category_code(self, user_data: UserData):
666        base_url = await self.base_url
667        header_clientid = self._get_header_client_id()
668        product_request = PreparedRequest(base_url, self.session, {"header_clientid": header_clientid})
669        response = await product_request.request(
670            "get",
671            "api/v1/plugins?apiLevel=99999&type=2",
672            headers={
673                "Authorization": user_data.token,
674            },
675        )
676        return {r["category"]: r["url"] for r in response["data"]["categoryPluginList"]}
677
678
679class PreparedRequest:
680    def __init__(
681        self, base_url: str, session: aiohttp.ClientSession | None = None, base_headers: dict | None = None
682    ) -> None:
683        self.base_url = base_url
684        self.base_headers = base_headers or {}
685        self.session = session
686
687    async def request(self, method: str, url: str, params=None, data=None, headers=None, json=None) -> dict:
688        _url = "/".join(s.strip("/") for s in [self.base_url, url])
689        _headers = {**self.base_headers, **(headers or {})}
690        close_session = self.session is None
691        session = self.session if self.session is not None else aiohttp.ClientSession()
692        try:
693            async with session.request(method, _url, params=params, data=data, headers=_headers, json=json) as resp:
694                return await resp.json()
695        except ContentTypeError as err:
696            """If we get an error, lets log everything for debugging."""
697            try:
698                resp_json = await resp.json(content_type=None)
699                _LOGGER.info("Resp: %s", resp_json)
700            except ContentTypeError as err_2:
701                _LOGGER.info(err_2)
702            resp_raw = await resp.read()
703            _LOGGER.info("Resp raw: %s", resp_raw)
704            # Still raise the err so that it's clear it failed.
705            raise err
706        finally:
707            if close_session:
708                await session.close()
709
710
711def _process_extra_hawk_values(values: dict | None) -> str:
712    if values is None:
713        return ""
714    else:
715        sorted_keys = sorted(values.keys())
716        result = []
717        for key in sorted_keys:
718            value = values.get(key)
719            result.append(f"{key}={value}")
720        return hashlib.md5("&".join(result).encode()).hexdigest()
721
722
723def _get_hawk_authentication(rriot: RRiot, url: str, formdata: dict | None = None, params: dict | None = None) -> str:
724    timestamp = math.floor(time.time())
725    nonce = secrets.token_urlsafe(6)
726    formdata_str = _process_extra_hawk_values(formdata)
727    params_str = _process_extra_hawk_values(params)
728
729    prestr = ":".join(
730        [
731            rriot.u,
732            rriot.s,
733            nonce,
734            str(timestamp),
735            hashlib.md5(url.encode()).hexdigest(),
736            params_str,
737            formdata_str,
738        ]
739    )
740    mac = base64.b64encode(hmac.new(rriot.h.encode(), prestr.encode(), hashlib.sha256).digest()).decode()
741    return f'Hawk id="{rriot.u}",s="{rriot.s}",ts="{timestamp}",nonce="{nonce}",mac="{mac}"'
742
743
744class UserWebApiClient:
745    """Wrapper around RoborockApiClient to provide information for a specific user.
746
747    This binds a RoborockApiClient to a specific user context with the
748    provided UserData. This allows for easier access to user-specific data,
749    to avoid needing to pass UserData around and mock out the web API.
750    """
751
752    def __init__(self, web_api: RoborockApiClient, user_data: UserData) -> None:
753        """Initialize the wrapper with the API client and user data."""
754        self._web_api = web_api
755        self._user_data = user_data
756
757    async def get_home_data(self) -> HomeData:
758        """Fetch home data using the API client."""
759        return await self._web_api.get_home_data_v3(self._user_data)
760
761    async def get_routines(self, device_id: str) -> list[HomeDataScene]:
762        """Fetch routines (scenes) for a specific device."""
763        return await self._web_api.get_scenes(self._user_data, device_id)
764
765    async def execute_routine(self, scene_id: int) -> None:
766        """Execute a specific routine (scene) by its ID."""
767        await self._web_api.execute_scene(self._user_data, scene_id)