diff --git a/src/apron_tools/providers/google/gmail/tools.py b/src/apron_tools/providers/google/gmail/tools.py index e9914e0..ca1cf58 100644 --- a/src/apron_tools/providers/google/gmail/tools.py +++ b/src/apron_tools/providers/google/gmail/tools.py @@ -3,7 +3,6 @@ from __future__ import annotations import base64 -import binascii import contextlib from email.mime.text import MIMEText @@ -58,6 +57,29 @@ def _headers(token: str, *, content_type: bool = False) -> dict[str, str]: return h +def _decode_base64url(data: str) -> bytes: + """Decode Gmail base64url data, restoring the padding Gmail omits. + + The Gmail API returns ``MessagePartBody.data`` and attachment data as + base64url that commonly drops the trailing ``=``, which + ``base64.urlsafe_b64decode`` rejects. Restoring it is a no-op on input + that is already padded. Decoding is strict: characters outside the + base64url alphabet are rejected rather than silently discarded, so + corrupt data surfaces as an error instead of truncated bytes. + + Args: + data: The base64url-encoded string to decode, with or without padding. + + Returns: + The decoded bytes. + + Raises: + ValueError: If ``data`` cannot be decoded, such as an invalid length, + a character outside the base64url alphabet, or non-ASCII input. + """ + return base64.b64decode(data + "=" * (-len(data) % 4), altchars=b"-_", validate=True) + + def _extract_header(headers: list[dict[str, str]], name: str) -> str: """Extract a single header value from a Gmail payload headers list.""" for h in headers: @@ -76,8 +98,8 @@ def _extract_body(payload: dict) -> str: body_data = payload.get("body", {}).get("data", "") if body_data: try: - return base64.urlsafe_b64decode(body_data).decode("utf-8") - except Exception: + return _decode_base64url(body_data).decode("utf-8") + except ValueError: return "(Could not decode email body)" parts = payload.get("parts", []) @@ -89,11 +111,11 @@ def _extract_body(payload: dict) -> str: part_data = part.get("body", {}).get("data", "") if mime_type == "text/plain" and part_data: - with contextlib.suppress(Exception): - plain_text = base64.urlsafe_b64decode(part_data).decode("utf-8") + with contextlib.suppress(ValueError): + plain_text = _decode_base64url(part_data).decode("utf-8") elif mime_type == "text/html" and part_data: - with contextlib.suppress(Exception): - html_text = base64.urlsafe_b64decode(part_data).decode("utf-8") + with contextlib.suppress(ValueError): + html_text = _decode_base64url(part_data).decode("utf-8") elif part.get("parts"): nested = _extract_body(part) if nested and nested != "(No email body found)": @@ -292,12 +314,9 @@ async def gmail_get_attachment( if not encoded: return GetAttachmentResult(success=False, error="Attachment contained no data.") - # Gmail returns base64url data that often omits trailing padding, which - # base64.urlsafe_b64decode rejects; restore it before decoding. - encoded += "=" * (-len(encoded) % 4) try: - raw = base64.urlsafe_b64decode(encoded) - except (binascii.Error, ValueError): + raw = _decode_base64url(encoded) + except ValueError: return GetAttachmentResult(success=False, error="Could not decode attachment data.") return GetAttachmentResult( diff --git a/tests/providers/google/gmail/test_tools.py b/tests/providers/google/gmail/test_tools.py index 9f9a622..61a7051 100644 --- a/tests/providers/google/gmail/test_tools.py +++ b/tests/providers/google/gmail/test_tools.py @@ -195,6 +195,172 @@ async def test_encodes_message_id_in_request_path(self, httpx_mock: HTTPXMock) - assert result.success is True + async def test_success_with_unpadded_single_part_body(self, httpx_mock: HTTPXMock) -> None: + # Gmail commonly returns MessagePartBody.data without trailing '=' padding. + body = "Hi Bob, following up on the thread." + padded = base64.urlsafe_b64encode(body.encode()).decode("ascii") + assert padded.endswith("="), "test payload must exercise padding removal" + + message = _load_json("get_message_full.json") + message["payload"]["mimeType"] = "text/plain" + message["payload"]["body"] = {"size": len(body), "data": padded.rstrip("=")} + message["payload"].pop("parts") + + httpx_mock.add_response( + url=f"{_GMAIL_BASE}/messages/msg-001?format=full", + json=message, + ) + + result = await gmail_read_email( + ReadEmailParams(message_id="msg-001"), + token=_TOKEN, + ) + + assert result.success is True + assert result.body == body + + async def test_success_with_unpadded_multipart_body(self, httpx_mock: HTTPXMock) -> None: + body = "Hi Bob, following up on the thread." + padded = base64.urlsafe_b64encode(body.encode()).decode("ascii") + assert padded.endswith("="), "test payload must exercise padding removal" + + message = _load_json("get_message_full.json") + message["payload"]["parts"][0]["body"] = { + "size": len(body), + "data": padded.rstrip("="), + } + + httpx_mock.add_response( + url=f"{_GMAIL_BASE}/messages/msg-001?format=full", + json=message, + ) + + result = await gmail_read_email( + ReadEmailParams(message_id="msg-001"), + token=_TOKEN, + ) + + assert result.success is True + assert result.body == body + + async def test_success_with_unpadded_html_body(self, httpx_mock: HTTPXMock) -> None: + # Unlike the other unpadded regression tests above, this payload requires + # two padding characters ('==') rather than just one. + body = "

Hi Bob, following up.

" + padded = base64.urlsafe_b64encode(body.encode()).decode("ascii") + assert padded.endswith("=="), "test payload must require two padding characters" + + message = _load_json("get_message_full.json") + message["payload"]["parts"][0] = { + "mimeType": "text/html", + "body": {"size": len(body), "data": padded.rstrip("=")}, + } + + httpx_mock.add_response( + url=f"{_GMAIL_BASE}/messages/msg-001?format=full", + json=message, + ) + + result = await gmail_read_email( + ReadEmailParams(message_id="msg-001"), + token=_TOKEN, + ) + + assert result.success is True + assert result.body == body + + async def test_single_part_body_with_undecodable_base64_returns_placeholder(self, httpx_mock: HTTPXMock) -> None: + # An invalid base64url length (a single data character) raises + # binascii.Error, a ValueError subclass, which the single-part path + # catches and reports as a decode failure. + message = _load_json("get_message_full.json") + message["payload"]["mimeType"] = "text/plain" + message["payload"]["body"] = {"size": 1, "data": "A"} + message["payload"].pop("parts") + + httpx_mock.add_response( + url=f"{_GMAIL_BASE}/messages/msg-001?format=full", + json=message, + ) + + result = await gmail_read_email( + ReadEmailParams(message_id="msg-001"), + token=_TOKEN, + ) + + assert result.success is True + assert result.body == "(Could not decode email body)" + + async def test_single_part_body_with_invalid_base64_alphabet_returns_placeholder( + self, httpx_mock: HTTPXMock + ) -> None: + # A complete valid Base64 block ("YWJj" -> "abc") followed by a stray + # character outside the base64url alphabet. The lenient decoder would + # discard the stray character and still decode "abc"; strict validation + # rejects it. This isolates the alphabet check rather than tripping the + # padding check the way a shorter invalid payload would. + message = _load_json("get_message_full.json") + message["payload"]["mimeType"] = "text/plain" + message["payload"]["body"] = {"size": 3, "data": "YWJj*"} + message["payload"].pop("parts") + + httpx_mock.add_response( + url=f"{_GMAIL_BASE}/messages/msg-001?format=full", + json=message, + ) + + result = await gmail_read_email( + ReadEmailParams(message_id="msg-001"), + token=_TOKEN, + ) + + assert result.success is True + assert result.body == "(Could not decode email body)" + + async def test_single_part_body_with_non_utf8_bytes_returns_placeholder(self, httpx_mock: HTTPXMock) -> None: + # Valid base64url that decodes to bytes that are not valid UTF-8 raises + # UnicodeDecodeError, also a ValueError subclass, on the single-part path. + non_utf8 = base64.urlsafe_b64encode(b"\xff\xfe").decode("ascii") + + message = _load_json("get_message_full.json") + message["payload"]["mimeType"] = "text/plain" + message["payload"]["body"] = {"size": 2, "data": non_utf8} + message["payload"].pop("parts") + + httpx_mock.add_response( + url=f"{_GMAIL_BASE}/messages/msg-001?format=full", + json=message, + ) + + result = await gmail_read_email( + ReadEmailParams(message_id="msg-001"), + token=_TOKEN, + ) + + assert result.success is True + assert result.body == "(Could not decode email body)" + + async def test_multipart_body_with_non_utf8_bytes_falls_through(self, httpx_mock: HTTPXMock) -> None: + # On the multipart path an undecodable part is suppressed, leaving no + # usable body, so extraction falls through to the sentinel. + non_utf8 = base64.urlsafe_b64encode(b"\xff\xfe").decode("ascii") + + message = _load_json("get_message_full.json") + message["payload"]["parts"][0]["body"] = {"size": 2, "data": non_utf8} + + httpx_mock.add_response( + url=f"{_GMAIL_BASE}/messages/msg-001?format=full", + json=message, + ) + + result = await gmail_read_email( + ReadEmailParams(message_id="msg-001"), + token=_TOKEN, + ) + + assert result.success is True + assert result.body == "(No email body found)" + async def test_api_error(self, httpx_mock: HTTPXMock) -> None: httpx_mock.add_response(status_code=404, text="Not Found") @@ -891,6 +1057,28 @@ async def test_success_with_unpadded_data(self, httpx_mock: HTTPXMock) -> None: assert result.data == raw assert result.size == len(raw) + async def test_success_with_unpadded_data_requiring_double_padding(self, httpx_mock: HTTPXMock) -> None: + # Unlike test_success_with_unpadded_data above, which only strips one + # padding character, this payload requires two ('=='). + raw = b"unpadded pdf report bytes" + padded = base64.urlsafe_b64encode(raw).decode("ascii") + assert padded.endswith("=="), "test payload must require two padding characters" + unpadded = padded.rstrip("=") + + httpx_mock.add_response( + url=f"{_GMAIL_BASE}/messages/msg-001/attachments/att-006", + json={"attachmentId": "att-006", "size": len(raw), "data": unpadded}, + ) + + result = await gmail_get_attachment( + GetAttachmentParams(message_id="msg-001", attachment_id="att-006"), + token=_TOKEN, + ) + + assert result.success is True + assert result.data == raw + assert result.size == len(raw) + async def test_non_json_response(self, httpx_mock: HTTPXMock) -> None: httpx_mock.add_response( url=f"{_GMAIL_BASE}/messages/msg-001/attachments/att-005",