149 lines
5.0 KiB
Python
149 lines
5.0 KiB
Python
"""Decode Graph fileAttachment contentBytes into PDF / DOC / DOCX files."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import base64
|
|
import binascii
|
|
import io
|
|
import zipfile
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
class AttachmentDecodeError(ValueError):
|
|
"""Raised when contentBytes is malformed or is not the expected format."""
|
|
|
|
|
|
_DEFAULT_OUT_DIR = Path(__file__).resolve().parent / "decoded_attachments"
|
|
|
|
|
|
def _decode_bytes(attachment: dict) -> bytes:
|
|
"""base64 -> raw bytes.
|
|
|
|
Graph's ``size`` often includes MIME/encoding overhead and may not equal
|
|
``len(contentBytes)`` after decode, so it is not treated as a hard check.
|
|
"""
|
|
b64 = attachment.get("contentBytes")
|
|
if not b64:
|
|
raise AttachmentDecodeError(f"{attachment.get('name')!r}: no contentBytes")
|
|
|
|
try:
|
|
return base64.b64decode(b64, validate=True)
|
|
except binascii.Error as exc:
|
|
raise AttachmentDecodeError(
|
|
f"{attachment.get('name')!r}: bad base64: {exc}"
|
|
) from exc
|
|
|
|
|
|
def _write(out_dir: Path, name: str, raw: bytes) -> Path:
|
|
out_dir = Path(out_dir)
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
dest = out_dir / Path(name).name # basename only — strip path traversal
|
|
dest.write_bytes(raw)
|
|
return dest
|
|
|
|
|
|
def decode_pdf(attachment: dict, out_dir: str | Path) -> Path:
|
|
"""Decode a PDF attachment and write it under out_dir."""
|
|
raw = _decode_bytes(attachment)
|
|
name = attachment.get("name")
|
|
if not raw.startswith(b"%PDF-"):
|
|
raise AttachmentDecodeError(f"{name!r}: not a PDF (missing %PDF- header)")
|
|
if b"%%EOF" not in raw[-2048:]:
|
|
raise AttachmentDecodeError(f"{name!r}: not a PDF (missing %%EOF trailer)")
|
|
return _write(Path(out_dir), name or "attachment.pdf", raw)
|
|
|
|
|
|
def decode_docx(attachment: dict, out_dir: str | Path) -> Path:
|
|
"""Decode a DOCX attachment and write it under out_dir."""
|
|
raw = _decode_bytes(attachment)
|
|
name = attachment.get("name")
|
|
if not raw.startswith(b"PK\x03\x04"):
|
|
raise AttachmentDecodeError(f"{name!r}: not a DOCX (missing ZIP signature)")
|
|
|
|
bio = io.BytesIO(raw)
|
|
if not zipfile.is_zipfile(bio):
|
|
raise AttachmentDecodeError(f"{name!r}: not a DOCX (invalid ZIP)")
|
|
bio.seek(0)
|
|
with zipfile.ZipFile(bio) as zf:
|
|
if not any(member.startswith("word/") for member in zf.namelist()):
|
|
raise AttachmentDecodeError(f"{name!r}: not a DOCX (no word/ entry)")
|
|
|
|
return _write(Path(out_dir), name or "attachment.docx", raw)
|
|
|
|
|
|
def decode_doc(attachment: dict, out_dir: str | Path) -> Path:
|
|
"""Decode a legacy DOC (OLE2) attachment and write it under out_dir."""
|
|
raw = _decode_bytes(attachment)
|
|
name = attachment.get("name")
|
|
if raw.startswith(b"PK\x03\x04"):
|
|
raise AttachmentDecodeError(
|
|
f"{name!r}: named .doc but content is DOCX — use decode_docx"
|
|
)
|
|
ole2 = b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1"
|
|
if not raw.startswith(ole2):
|
|
raise AttachmentDecodeError(f"{name!r}: not a DOC (missing OLE2 signature)")
|
|
return _write(Path(out_dir), name or "attachment.doc", raw)
|
|
|
|
|
|
_DECODERS = {
|
|
".pdf": decode_pdf,
|
|
".docx": decode_docx,
|
|
".doc": decode_doc,
|
|
}
|
|
|
|
|
|
def _decode_one(attachment: dict, out_dir: str | Path) -> Path:
|
|
"""Route on the file extension to the right decoder."""
|
|
ext = Path(attachment.get("name", "")).suffix.lower()
|
|
if ext not in _DECODERS:
|
|
raise AttachmentDecodeError(f"unsupported extension {ext!r}")
|
|
return _DECODERS[ext](attachment, out_dir)
|
|
|
|
|
|
def _normalize_attachments(attachments: Any) -> list[dict]:
|
|
"""Accept None, a single dict, or a list; return only dict items."""
|
|
if attachments is None:
|
|
return []
|
|
if isinstance(attachments, dict):
|
|
return [attachments]
|
|
if isinstance(attachments, list):
|
|
return [a for a in attachments if isinstance(a, dict)]
|
|
return []
|
|
|
|
|
|
def _decode_attachments_sync(
|
|
attachments: Any,
|
|
out_dir: str | Path | None = None,
|
|
) -> list[str]:
|
|
"""Decode supported file attachments; skip empty / non-file / unsupported."""
|
|
dest_dir = Path(out_dir) if out_dir is not None else _DEFAULT_OUT_DIR
|
|
paths: list[str] = []
|
|
|
|
for attachment in _normalize_attachments(attachments):
|
|
# Graph itemAttachment / referenceAttachment have no contentBytes
|
|
if not attachment.get("contentBytes"):
|
|
continue
|
|
ext = Path(attachment.get("name") or "").suffix.lower()
|
|
if ext not in _DECODERS:
|
|
continue
|
|
path = _decode_one(attachment, dest_dir)
|
|
paths.append(str(path))
|
|
|
|
return paths
|
|
|
|
|
|
async def decode_attachment(
|
|
attachments: Any,
|
|
out_dir: str | Path | None = None,
|
|
) -> list[str]:
|
|
"""
|
|
Decode Graph attachments into files under out_dir.
|
|
|
|
Designed for views: ``await decode_attachment(data.get("attachments"))``.
|
|
Accepts None, a single attachment dict, or a list of attachment dicts.
|
|
Returns absolute/relative path strings for successfully written files.
|
|
"""
|
|
return await asyncio.to_thread(_decode_attachments_sync, attachments, out_dir)
|