diff --git a/gpmc/api.py b/gpmc/api.py index 461c1c5..474464a 100644 --- a/gpmc/api.py +++ b/gpmc/api.py @@ -2,13 +2,17 @@ import time from urllib.parse import parse_qs from pathlib import Path +import logging import requests from requests.adapters import HTTPAdapter, Retry from blackboxprotobuf import decode_message, encode_message +from blackboxprotobuf.lib.exceptions import BlackboxProtobufException from . import message_types -from .exceptions import UploadRejected +from .exceptions import UploadRejected, ProtobufDecodeError, ProtobufEncodeError + +logger = logging.getLogger("rich") DEFAULT_TIMEOUT = 60 RETRIES = 10 @@ -233,7 +237,12 @@ def upload_file(self, file: str | Path | bytes | IO[bytes] | Generator[bytes, No response.raise_for_status() - upload_response_decoded, _ = decode_message(response.content) + try: + upload_response_decoded, _ = decode_message(response.content) + except (BlackboxProtobufException, Exception) as e: + logger.warning(f"Failed to decode upload response: {e}") + logger.debug(f"Response content (first 500 bytes): {response.content[:500]}") + raise ProtobufDecodeError(f"Failed to decode upload response: {e}") from e return upload_response_decoded def commit_upload( @@ -331,7 +340,12 @@ def commit_upload( "3": bytes([1, 3]), } - serialized_data = encode_message(proto_body, message_types.COMMIT_UPLOAD) # type: ignore + try: + serialized_data = encode_message(proto_body, message_types.COMMIT_UPLOAD) # type: ignore + except (BlackboxProtobufException, Exception) as e: + logger.error(f"Failed to encode commit_upload request for file '{file_name}'") + logger.debug(f"Proto body keys: {proto_body.keys()}") + raise ProtobufEncodeError(f"Failed to encode commit_upload request: {e}") from e headers = { "Accept-Encoding": "gzip", @@ -350,10 +364,16 @@ def commit_upload( timeout=self.timeout, ) response.raise_for_status() - decoded_message, _ = decode_message(response.content) + try: + decoded_message, _ = decode_message(response.content) + except (BlackboxProtobufException, Exception) as e: + logger.error(f"Failed to decode commit_upload response for file '{file_name}'") + logger.debug(f"Response content (first 500 bytes): {response.content[:500]}") + raise ProtobufDecodeError(f"Failed to decode commit_upload response: {e}") from e try: media_key = decoded_message["1"]["3"]["1"] except KeyError as e: + logger.warning(f"File upload rejected by api for file '{file_name}'") raise UploadRejected("File upload rejected by api") from e return media_key @@ -422,7 +442,13 @@ def create_album(self, album_name: str, media_keys: Sequence[str]) -> str: "8": {"3": self.model, "4": self.make, "5": self.android_api_version}, } - serialized_data = encode_message(proto_body, message_types.CREATE_ALBUM) # type: ignore + try: + serialized_data = encode_message(proto_body, message_types.CREATE_ALBUM) # type: ignore + except (BlackboxProtobufException, Exception) as e: + logger.error(f"Failed to encode create_album request for album '{album_name}' with {len(media_keys)} media items") + logger.debug(f"Proto body: {proto_body}") + logger.debug(f"Media keys: {media_keys[:5]}..." if len(media_keys) > 5 else f"Media keys: {media_keys}") + raise ProtobufEncodeError(f"Failed to encode create_album request: {e}") from e headers = { "Accept-Encoding": "gzip", @@ -442,8 +468,15 @@ def create_album(self, album_name: str, media_keys: Sequence[str]) -> str: ) response.raise_for_status() - decoded_message, _ = decode_message(response.content) - return decoded_message["1"]["1"] + try: + decoded_message, _ = decode_message(response.content) + return decoded_message["1"]["1"] + except (BlackboxProtobufException, KeyError, Exception) as e: + logger.error(f"Failed to decode create_album response for album '{album_name}'") + logger.debug(f"Response content (first 500 bytes): {response.content[:500]}") + if isinstance(e, KeyError): + raise ProtobufDecodeError(f"Unexpected response structure when creating album: {e}") from e + raise ProtobufDecodeError(f"Failed to decode create_album response: {e}") from e def add_media_to_album(self, album_media_key: str, media_keys: Sequence[str]) -> dict: """Add media to an album. diff --git a/gpmc/client.py b/gpmc/client.py index 02c3491..a42ccf9 100644 --- a/gpmc/client.py +++ b/gpmc/client.py @@ -27,6 +27,7 @@ from . import utils from .hash_handler import calculate_sha1_hash, convert_sha1_hash from .db_update_parser import parse_db_update +from .exceptions import ProtobufDecodeError, ProtobufEncodeError # Make Ctrl+C work for cancelling threads signal.signal(signal.SIGINT, signal.SIG_DFL) @@ -476,8 +477,13 @@ def _upload_concurrently(self, path_hash_pairs: Mapping[Path, bytes | str], thre try: media_key_dict = future.result() uploaded_files = uploaded_files | media_key_dict + except (ProtobufDecodeError, ProtobufEncodeError) as e: + self.logger.error(f"Protobuf error uploading file {file[0]}: {e}") + self.logger.info(f"Skipping {file[0]} due to protobuf encoding/decoding issue") + upload_error_count += 1 + overall_progress.update(task_id=overall_task_id, description=f"[bold red] Errors: {upload_error_count}") except Exception as e: - self.logger.error(f"Error uploading file {file}: {e}") + self.logger.error(f"Error uploading file {file[0]}: {e}") upload_error_count += 1 overall_progress.update(task_id=overall_task_id, description=f"[bold red] Errors: {upload_error_count}") finally: @@ -562,13 +568,18 @@ def add_to_album(self, media_keys: Sequence[str], album_name: str, show_progress current_album_key = None for j in range(0, len(album_batch), batch_size): batch = album_batch[j : j + batch_size] - if current_album_key is None: - # Create the album with the first batch - current_album_key = self.api.create_album(album_name=current_album_name, media_keys=batch) - album_keys.append(current_album_key) - else: - # Add to the existing album - self.api.add_media_to_album(album_media_key=current_album_key, media_keys=batch) + try: + if current_album_key is None: + # Create the album with the first batch + current_album_key = self.api.create_album(album_name=current_album_name, media_keys=batch) + album_keys.append(current_album_key) + else: + # Add to the existing album + self.api.add_media_to_album(album_media_key=current_album_key, media_keys=batch) + except (ProtobufDecodeError, ProtobufEncodeError) as e: + self.logger.error(f"Protobuf error while adding items to album '{current_album_name}': {e}") + self.logger.info(f"Batch of {len(batch)} items skipped. Media keys (first 3): {batch[:3]}") + raise progress.update(task, advance=len(batch)) album_counter += 1 return album_keys diff --git a/gpmc/exceptions.py b/gpmc/exceptions.py index 8ab022d..100f6d2 100644 --- a/gpmc/exceptions.py +++ b/gpmc/exceptions.py @@ -4,3 +4,13 @@ class CustomException(Exception): class UploadRejected(CustomException): pass + + +class ProtobufDecodeError(CustomException): + """Raised when protobuf message decoding fails.""" + pass + + +class ProtobufEncodeError(CustomException): + """Raised when protobuf message encoding fails.""" + pass