pypetkitapi 1.9.2__py3-none-any.whl → 1.9.4__py3-none-any.whl

pypetkitapi/__init__.py

@@ -36,11 +36,11 @@ from .containers import Pet
 from .exceptions import PetkitAuthenticationError, PypetkitError
 from .feeder_container import Feeder, RecordsItems
 from .litter_container import Litter, LitterRecord, WorkState
-from .medias import MediaHandler, MediasFiles
+from .media import DownloadDecryptMedia, MediaFile, MediaManager
 from .purifier_container import Purifier
 from .water_fountain_container import WaterFountain
 
-__version__ = "1.9.2"
+__version__ = "1.9.4"
 
 __all__ = [
     "CTW3",
@@ -65,8 +65,9 @@ __all__ = [
     "Litter",
     "LitterCommand",
     "LitterRecord",
-    "MediaHandler",
-    "MediasFiles",
+    "MediaManager",
+    "DownloadDecryptMedia",
+    "MediaFile",
     "Pet",
     "PetCommand",
     "PetKitClient",

pypetkitapi/client.py

@@ -11,12 +11,14 @@ import urllib.parse
 
 import aiohttp
 from aiohttp import ContentTypeError
+import m3u8
 
 from pypetkitapi.command import ACTIONS_MAP, FOUNTAIN_COMMAND, FountainAction
 from pypetkitapi.const import (
     BLE_CONNECT_ATTEMPT,
     BLE_END_TRAME,
     BLE_START_TRAME,
+    CLIENT_NFO,
     DEVICE_DATA,
     DEVICE_RECORDS,
     DEVICE_STATS,
@@ -25,6 +27,8 @@ from pypetkitapi.const import (
     DEVICES_PURIFIER,
     DEVICES_WATER_FOUNTAIN,
     ERR_KEY,
+    LITTER_NO_CAMERA,
+    LITTER_WITH_CAMERA,
     LOGIN_DATA,
     PET,
     RES_KEY,
@@ -50,6 +54,7 @@ from pypetkitapi.exceptions import (
     PetkitInvalidHTTPResponseCodeError,
     PetkitInvalidResponseFormat,
     PetkitRegionalServerNotFoundError,
+    PetkitSessionError,
     PetkitSessionExpiredError,
     PetkitTimeoutError,
     PypetkitError,
@@ -57,8 +62,22 @@ from pypetkitapi.exceptions import (
 from pypetkitapi.feeder_container import Feeder, FeederRecord
 from pypetkitapi.litter_container import Litter, LitterRecord, LitterStats, PetOutGraph
 from pypetkitapi.purifier_container import Purifier
+from pypetkitapi.utils import get_timezone_offset
 from pypetkitapi.water_fountain_container import WaterFountain, WaterFountainRecord
 
+data_handlers = {}
+
+
+def data_handler(data_type):
+    """Register a data handler for a specific data type."""
+
+    def wrapper(func):
+        data_handlers[data_type] = func
+        return func
+
+    return wrapper
+
+
 _LOGGER = logging.getLogger(__name__)
 
 
@@ -86,7 +105,9 @@ class PetKitClient:
         self.petkit_entities = {}
         self.aiohttp_session = session or aiohttp.ClientSession()
         self.req = PrepReq(
-            base_url=PetkitDomain.PASSPORT_PETKIT, session=self.aiohttp_session
+            base_url=PetkitDomain.PASSPORT_PETKIT,
+            session=self.aiohttp_session,
+            timezone=self.timezone,
         )
 
     async def _get_base_url(self) -> None:
@@ -129,12 +150,18 @@ class PetKitClient:
     async def login(self, valid_code: str | None = None) -> None:
         """Login to the PetKit service and retrieve the appropriate server."""
         # Retrieve the list of servers
+        self._session = None
         await self._get_base_url()
 
         _LOGGER.info("Logging in to PetKit server")
 
         # Prepare the data to send
+        client_nfo = CLIENT_NFO.copy()
+        client_nfo["timezoneId"] = self.timezone
+        client_nfo["timezone"] = get_timezone_offset(self.timezone)
+
         data = LOGIN_DATA.copy()
+        data["client"] = str(client_nfo)
         data["encrypt"] = "1"
         data["region"] = self.region
         data["username"] = self.username
@@ -156,6 +183,8 @@ class PetKitClient:
         )
         session_data = response["session"]
         self._session = SessionInfo(**session_data)
+        expiration_date = datetime.now() + timedelta(seconds=self._session.expires_in)
+        _LOGGER.debug("Login successful (token expiration %s)", expiration_date)
 
     async def refresh_session(self) -> None:
         """Refresh the session."""
@@ -169,6 +198,7 @@ class PetKitClient:
         session_data = response["session"]
         self._session = SessionInfo(**session_data)
         self._session.refreshed_at = datetime.now().strftime("%Y-%m-%dT%H:%M:%S.%f")
+        _LOGGER.debug("Session refreshed at %s", self._session.refreshed_at)
 
     async def validate_session(self) -> None:
         """Check if the session is still valid and refresh or re-login if necessary."""
@@ -177,31 +207,27 @@ class PetKitClient:
             await self.login()
             return
 
-        created_at = datetime.strptime(
-            self._session.created_at,
-            "%Y-%m-%dT%H:%M:%S.%f%z",
+        created = datetime.strptime(self._session.created_at, "%Y-%m-%dT%H:%M:%S.%f%z")
+        is_expired = datetime.now(tz=created.tzinfo) - created >= timedelta(
+            seconds=self._session.expires_in
         )
-        current_time = datetime.now(tz=created_at.tzinfo)
-        token_age = current_time - created_at
-        max_age = timedelta(seconds=self._session.expires_in)
-        half_max_age = max_age / 2
 
-        if token_age > max_age:
+        if is_expired:
             _LOGGER.debug("Token expired, re-logging in")
             await self.login()
-        elif half_max_age < token_age <= max_age:
-            _LOGGER.debug("Token still OK, but refreshing session")
-            await self.refresh_session()
+        # elif (max_age / 2) < token_age < max_age:
+        #     _LOGGER.debug("Token still OK, but refreshing session")
+        #     await self.refresh_session()
 
     async def get_session_id(self) -> dict:
         """Return the session ID."""
+        await self.validate_session()
         if self._session is None:
-            raise PypetkitError("Session is not initialized.")
+            raise PetkitSessionError("No session ID available")
         return {"F-Session": self._session.id, "X-Session": self._session.id}
 
     async def _get_account_data(self) -> None:
         """Get the account data from the PetKit service."""
-        await self.validate_session()
         _LOGGER.debug("Fetching account data")
         response = await self.req.request(
             method=HTTPMethod.GET,
@@ -223,27 +249,40 @@ class PetKitClient:
                         groupId=0,
                         type=0,
                         typeCode=0,
-                        uniqueId=pet.sn,
+                        uniqueId=str(pet.sn),
                     )
 
     async def get_devices_data(self) -> None:
         """Get the devices data from the PetKit servers."""
-        await self.validate_session()
-
         start_time = datetime.now()
         if not self.account_data:
             await self._get_account_data()
 
-        main_tasks = []
-        record_tasks = []
-        device_list: list[Device] = []
+        device_list = self._collect_devices()
+        main_tasks, record_tasks = self._prepare_tasks(device_list)
 
+        await asyncio.gather(*main_tasks)
+        await asyncio.gather(*record_tasks)
+        await self._execute_stats_tasks()
+
+        end_time = datetime.now()
+        _LOGGER.debug("Petkit data fetched successfully in: %s", end_time - start_time)
+
+    def _collect_devices(self) -> list[Device]:
+        """Collect all devices from account data."""
+        device_list = []
         for account in self.account_data:
             _LOGGER.debug("List devices data for account: %s", account)
             if account.device_list:
                 _LOGGER.debug("Devices in account: %s", account.device_list)
                 device_list.extend(account.device_list)
                 _LOGGER.debug("Found %s devices", len(account.device_list))
+        return device_list
+
+    def _prepare_tasks(self, device_list: list[Device]) -> tuple[list, list]:
+        """Prepare main and record tasks based on device types."""
+        main_tasks = []
+        record_tasks = []
 
         for device in device_list:
             device_type = device.device_type
@@ -253,15 +292,9 @@ class PetKitClient:
                 record_tasks.append(self._fetch_device_data(device, FeederRecord))
 
             elif device_type in DEVICES_LITTER_BOX:
-                main_tasks.append(
-                    self._fetch_device_data(device, Litter),
-                )
+                main_tasks.append(self._fetch_device_data(device, Litter))
                 record_tasks.append(self._fetch_device_data(device, LitterRecord))
-
-                if device_type in [T3, T4]:
-                    record_tasks.append(self._fetch_device_data(device, LitterStats))
-                if device_type in [T5, T6]:
-                    record_tasks.append(self._fetch_device_data(device, PetOutGraph))
+                self._add_litter_box_tasks(record_tasks, device_type, device)
 
             elif device_type in DEVICES_WATER_FOUNTAIN:
                 main_tasks.append(self._fetch_device_data(device, WaterFountain))
@@ -272,26 +305,26 @@ class PetKitClient:
             elif device_type in DEVICES_PURIFIER:
                 main_tasks.append(self._fetch_device_data(device, Purifier))
 
-        # Execute main device tasks first
-        await asyncio.gather(*main_tasks)
+        return main_tasks, record_tasks
 
-        # Then execute record tasks
-        await asyncio.gather(*record_tasks)
+    def _add_litter_box_tasks(
+        self, record_tasks: list, device_type: str, device: Device
+    ):
+        """Add specific tasks for litter box devices."""
+        if device_type in [T3, T4]:
+            record_tasks.append(self._fetch_device_data(device, LitterStats))
+        if device_type in [T5, T6]:
+            record_tasks.append(self._fetch_device_data(device, PetOutGraph))
 
-        # Add populate_pet_stats tasks
+    async def _execute_stats_tasks(self) -> None:
+        """Execute tasks to populate pet stats."""
         stats_tasks = [
-            self.populate_pet_stats(self.petkit_entities[device.device_id])
-            for device in device_list
-            if device.device_type in DEVICES_LITTER_BOX
+            self.populate_pet_stats(entity)
+            for device_id, entity in self.petkit_entities.items()
+            if isinstance(entity, Litter)
         ]
-
-        # Execute stats tasks
         await asyncio.gather(*stats_tasks)
 
-        end_time = datetime.now()
-        total_time = end_time - start_time
-        _LOGGER.debug("Petkit data fetched successfully in: %s", total_time)
-
     async def _fetch_device_data(
         self,
         device: Device,
@@ -345,23 +378,48 @@ class PetKitClient:
             _LOGGER.error("Unexpected response type: %s", type(response))
             return
 
-        if data_class.data_type == DEVICE_DATA:
-            self.petkit_entities[device.device_id] = device_data
-            self.petkit_entities[device.device_id].device_nfo = device
-            _LOGGER.debug("Device data fetched OK for %s", device_type)
-        elif data_class.data_type == DEVICE_RECORDS:
-            self.petkit_entities[device.device_id].device_records = device_data
+        # Dispatch to the appropriate handler
+        handler = data_handlers.get(data_class.data_type)
+        if handler:
+            await handler(self, device, device_data, device_type)
+        else:
+            _LOGGER.error("Unknown data type: %s", data_class.data_type)
+
+    @data_handler(DEVICE_DATA)
+    async def _handle_device_data(self, device, device_data, device_type):
+        """Handle device data."""
+        self.petkit_entities[device.device_id] = device_data
+        self.petkit_entities[device.device_id].device_nfo = device
+        _LOGGER.debug("Device data fetched OK for %s", device_type)
+
+    @data_handler(DEVICE_RECORDS)
+    async def _handle_device_records(self, device, device_data, device_type):
+        """Handle device records."""
+        entity = self.petkit_entities.get(device.device_id)
+        if entity and isinstance(entity, (Feeder, Litter, WaterFountain)):
+            entity.device_records = device_data
             _LOGGER.debug("Device records fetched OK for %s", device_type)
-        elif data_class.data_type == DEVICE_STATS:
-            if device_type in [T3, T4]:
-                self.petkit_entities[device.device_id].device_stats = device_data
-            if device_type in [T5, T6]:
-                self.petkit_entities[device.device_id].device_pet_graph_out = (
-                    device_data
-                )
+        else:
+            _LOGGER.warning(
+                "Cannot assign device_records to entity of type %s",
+                type(entity),
+            )
+
+    @data_handler(DEVICE_STATS)
+    async def _handle_device_stats(self, device, device_data, device_type):
+        """Handle device stats."""
+        entity = self.petkit_entities.get(device.device_id)
+        if isinstance(entity, Litter):
+            if device_type in LITTER_NO_CAMERA:
+                entity.device_stats = device_data
+            if device_type in LITTER_WITH_CAMERA:
+                entity.device_pet_graph_out = device_data
             _LOGGER.debug("Device stats fetched OK for %s", device_type)
         else:
-            _LOGGER.error("Unknown data type: %s", data_class.data_type)
+            _LOGGER.warning(
+                "Cannot assign device_stats or device_pet_graph_out to entity of type %s",
+                type(entity),
+            )
 
     async def get_pets_list(self) -> list[Pet]:
         """Extract and return the list of pets."""
@@ -386,6 +444,12 @@ class PetKitClient:
     async def populate_pet_stats(self, stats_data: Litter) -> None:
         """Collect data from litter data to populate pet stats."""
 
+        if not stats_data.device_nfo:
+            _LOGGER.warning(
+                "No device info for %s can't populate pet infos", stats_data
+            )
+            return
+
         pets_list = await self.get_pets_list()
         for pet in pets_list:
             if (
@@ -433,25 +497,27 @@ class PetKitClient:
     async def _get_fountain_instance(self, fountain_id: int) -> WaterFountain:
         # Retrieve the water fountain object
         water_fountain = self.petkit_entities.get(fountain_id)
-        if not water_fountain:
+        if not isinstance(water_fountain, WaterFountain):
             _LOGGER.error("Water fountain with ID %s not found.", fountain_id)
-            raise ValueError(f"Water fountain with ID {fountain_id} not found.")
+            raise TypeError(f"Water fountain with ID {fountain_id} not found.")
         return water_fountain
 
     async def check_relay_availability(self, fountain_id: int) -> bool:
         """Check if BLE relay is available for the account."""
         fountain = None
+
         for account in self.account_data:
-            fountain = next(
-                (
-                    device
-                    for device in account.device_list
-                    if device.device_id == fountain_id
-                ),
-                None,
-            )
-            if fountain:
-                break
+            if account.device_list:
+                fountain = next(
+                    (
+                        device
+                        for device in account.device_list
+                        if device.device_id == fountain_id
+                    ),
+                    None,
+                )
+                if fountain:
+                    break
 
         if not fountain:
             raise ValueError(
@@ -568,13 +634,13 @@ class PetKitClient:
             _LOGGER.error("BLE connection not established.")
             return False
 
-        command = FOUNTAIN_COMMAND.get[command, None]
-        if command is None:
+        command_data = FOUNTAIN_COMMAND.get(command)
+        if command_data is None:
             _LOGGER.error("Command not found.")
             return False
 
         cmd_code, cmd_data = await self.get_ble_cmd_data(
-            command, water_fountain.ble_counter
+            list(command_data), water_fountain.ble_counter
         )
 
         response = await self.req.request(
@@ -595,6 +661,46 @@ class PetKitClient:
         _LOGGER.info("BLE command sent successfully.")
         return True
 
+    async def get_cloud_video(self, video_url: str) -> dict[str, str | int]:
+        """Get the video m3u8 link from the cloud."""
+        response = await self.req.request(
+            method=HTTPMethod.POST,
+            url=video_url,
+            headers=await self.get_session_id(),
+        )
+        return response[0]
+
+    async def extract_segments_m3u8(self, m3u8_url: str) -> tuple[str, str, list[str]]:
+        """Extract segments from the m3u8 file.
+        :param: m3u8_url: URL of the m3u8 file
+        :return: aes_key, key_iv, segment_lst
+        """
+        # Extract segments from m3u8 file
+        response = await self.req.request(
+            method=HTTPMethod.GET,
+            url=m3u8_url,
+            headers=await self.get_session_id(),
+        )
+        m3u8_obj = m3u8.loads(response[RES_KEY])
+
+        if not m3u8_obj.segments or not m3u8_obj.keys:
+            raise PetkitInvalidResponseFormat("No segments or key found in m3u8 file.")
+
+        # Extract segments from m3u8 file
+        segment_lst = [segment.uri for segment in m3u8_obj.segments]
+        # Extract key_uri and key_iv from m3u8 file
+        key_uri = m3u8_obj.keys[0].uri
+        key_iv = str(m3u8_obj.keys[0].iv)
+
+        # Extract aes_key from video segments
+        response = await self.req.request(
+            method=HTTPMethod.GET,
+            url=key_uri,
+            full_url=True,
+            headers=await self.get_session_id(),
+        )
+        return response[RES_KEY], key_iv, segment_lst
+
     async def send_api_request(
         self,
         device_id: int,
@@ -602,11 +708,11 @@ class PetKitClient:
         setting: dict | None = None,
     ) -> bool:
         """Control the device using the PetKit API."""
-        await self.validate_session()
-
         device = self.petkit_entities.get(device_id, None)
         if not device:
             raise PypetkitError(f"Device with ID {device_id} not found.")
+        if device.device_nfo is None:
+            raise PypetkitError(f"Device with ID {device_id} has no device_nfo.")
 
         _LOGGER.debug(
             "Control API device=%s id=%s action=%s param=%s",
@@ -669,41 +775,43 @@ class PetKitClient:
 class PrepReq:
     """Prepare the request to the PetKit API."""
 
-    def __init__(self, base_url: str, session: aiohttp.ClientSession) -> None:
+    def __init__(
+        self, base_url: str, session: aiohttp.ClientSession, timezone: str
+    ) -> None:
         """Initialize the request."""
         self.base_url = base_url
         self.session = session
+        self.timezone = timezone
         self.base_headers = self._generate_header()
 
-    @staticmethod
-    def _generate_header() -> dict[str, str]:
+    def _generate_header(self) -> dict[str, str]:
         """Create header for interaction with API endpoint."""
-
         return {
             "Accept": Header.ACCEPT.value,
-            "Accept-Language": Header.ACCEPT_LANG,
-            "Accept-Encoding": Header.ENCODING,
-            "Content-Type": Header.CONTENT_TYPE,
-            "User-Agent": Header.AGENT,
-            "X-Img-Version": Header.IMG_VERSION,
-            "X-Locale": Header.LOCALE,
-            "X-Client": Header.CLIENT,
-            "X-Hour": Header.HOUR,
-            "X-TimezoneId": Header.TIMEZONE_ID,
-            "X-Api-Version": Header.API_VERSION,
-            "X-Timezone": Header.TIMEZONE,
+            "Accept-Language": Header.ACCEPT_LANG.value,
+            "Accept-Encoding": Header.ENCODING.value,
+            "Content-Type": Header.CONTENT_TYPE.value,
+            "User-Agent": Header.AGENT.value,
+            "X-Img-Version": Header.IMG_VERSION.value,
+            "X-Locale": Header.LOCALE.value,
+            "X-Client": Header.CLIENT.value,
+            "X-Hour": Header.HOUR.value,
+            "X-TimezoneId": self.timezone,
+            "X-Api-Version": Header.API_VERSION.value,
+            "X-Timezone": get_timezone_offset(self.timezone),
         }
 
     async def request(
         self,
         method: str,
         url: str,
+        full_url: bool = False,
         params=None,
         data=None,
         headers=None,
     ) -> dict:
         """Make a request to the PetKit API."""
-        _url = "/".join(s.strip("/") for s in [self.base_url, url])
+        _url = url if full_url else "/".join(s.strip("/") for s in [self.base_url, url])
         _headers = {**self.base_headers, **(headers or {})}
         _LOGGER.debug("Request: %s %s", method, _url)
         try:
@@ -729,12 +837,14 @@ class PrepReq:
             ) from e
 
         try:
-            response_json = await response.json()
+            if response.content_type == "application/json":
+                response_json = await response.json()
+            else:
+                return {RES_KEY: await response.text()}
         except ContentTypeError:
             raise PetkitInvalidResponseFormat(
                 "Response is not in JSON format"
             ) from None
-
         # Check for errors in the response
         if ERR_KEY in response_json:
             error_code = int(response_json[ERR_KEY].get("code", 0))

pypetkitapi/command.py

@@ -9,7 +9,6 @@ import json
 from pypetkitapi.const import (
     ALL_DEVICES,
     D3,
-    D4,
     D4H,
     D4S,
     D4SH,
@@ -145,25 +144,28 @@ class CmdData:
 
 def get_endpoint_manual_feed(device):
     """Get the endpoint for the device"""
-    if device.device_nfo.device_type == FEEDER_MINI:
-        return PetkitEndpoint.MANUAL_FEED_MINI
-    if device.device_nfo.device_type == FEEDER:
-        return PetkitEndpoint.MANUAL_FEED_FRESH_ELEMENT
-    return PetkitEndpoint.MANUAL_FEED_DUAL
+    if device.device_nfo.device_type in [FEEDER_MINI, FEEDER]:
+        return PetkitEndpoint.MANUAL_FEED_OLD  # Old endpoint snakecase
+    return PetkitEndpoint.MANUAL_FEED_NEW  # New endpoint camelcase
 
 
 def get_endpoint_reset_desiccant(device):
+    """Get the endpoint for the device"""
+    if device.device_nfo.device_type in [FEEDER_MINI, FEEDER]:
+        return PetkitEndpoint.DESICCANT_RESET_OLD  # Old endpoint snakecase
+    return PetkitEndpoint.DESICCANT_RESET_NEW  # New endpoint camelcase
+
+
+def get_endpoint_update_setting(device):
     """Get the endpoint for the device"""
     if device.device_nfo.device_type == FEEDER_MINI:
-        return PetkitEndpoint.MINI_DESICCANT_RESET
-    if device.device_nfo.device_type == FEEDER:
-        return PetkitEndpoint.FRESH_ELEMENT_DESICCANT_RESET
-    return PetkitEndpoint.DESICCANT_RESET
+        return PetkitEndpoint.UPDATE_SETTING_FEEDER_MINI
+    return PetkitEndpoint.UPDATE_SETTING
 
 
 ACTIONS_MAP = {
     DeviceCommand.UPDATE_SETTING: CmdData(
-        endpoint=PetkitEndpoint.UPDATE_SETTING,
+        endpoint=lambda device: get_endpoint_update_setting(device),
         params=lambda device, setting: {
             "id": device.id,
             "kv": json.dumps(setting),
@@ -199,16 +201,6 @@ ACTIONS_MAP = {
     ),
     FeederCommand.MANUAL_FEED: CmdData(
         endpoint=lambda device: get_endpoint_manual_feed(device),
-        params=lambda device, setting: {
-            "day": datetime.datetime.now().strftime("%Y%m%d"),
-            "deviceId": device.id,
-            "time": "-1",
-            **setting,
-        },
-        supported_device=[FEEDER, FEEDER_MINI, D3, D4, D4H],
-    ),
-    FeederCommand.MANUAL_FEED_DUAL: CmdData(
-        endpoint=PetkitEndpoint.MANUAL_FEED_DUAL,
         params=lambda device, setting: {
             "day": datetime.datetime.now().strftime("%Y%m%d"),
             "deviceId": device.id,
@@ -216,7 +208,7 @@ ACTIONS_MAP = {
             "time": "-1",
             **setting,
         },
-        supported_device=[D4S, D4SH],
+        supported_device=DEVICES_FEEDER,
     ),
     FeederCommand.CANCEL_MANUAL_FEED: CmdData(
         endpoint=lambda device: (