webmail-client/backend/tests/outbox/test_inline_images.py

"""Images quoted from another message go out inside the message, not as links
to this server; such links already sent load from here when shown."""

from __future__ import annotations

import email
import io
import shutil
import uuid
from collections.abc import AsyncIterator
from datetime import UTC, datetime
from email import policy
from email.message import EmailMessage
from pathlib import Path
from types import SimpleNamespace

import pytest_asyncio
from PIL import Image
from sqlalchemy import delete
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from sqlalchemy.pool import NullPool

from app.config import get_settings
from app.identity.passwords import hash_password
from app.indexer.indexer import persist_message
from app.indexer.sanitize import sanitize_html
from app.models import Account, Mailbox, User
from app.outbox import inline_images
from app.outbox.mime_build import build_mime
from app.providers.base import RemoteMessage


def _png() -> bytes:
    out = io.BytesIO()
    Image.new("RGB", (40, 20), (10, 120, 200)).save(out, "PNG")
    return out.getvalue()


LOGO = _png()


@pytest_asyncio.fixture
async def ctx() -> AsyncIterator[tuple[AsyncSession, uuid.UUID, uuid.UUID]]:
    engine = create_async_engine(get_settings().database_url, poolclass=NullPool)
    db = async_sessionmaker(engine, expire_on_commit=False)()
    user = User(handle=f"inl{uuid.uuid4().hex[:8]}@test.local", password_hash=hash_password("pw"))
    db.add(user)
    await db.flush()
    account = Account(
        user_id=user.id,
        email="me@example.com",
        imap_config={},
        smtp_config={},
        secret_ct=b"",
        discovery_src="manual",
    )
    db.add(account)
    await db.flush()
    inbox = Mailbox(account_id=account.id, user_id=user.id, path="INBOX", special_use="inbox")
    db.add(inbox)
    await db.commit()
    m = EmailMessage()
    m["Subject"] = "Signed"
    m["From"] = "Ben <ben@example.org>"
    m["To"] = "me@example.com"
    m["Message-ID"] = "<signed@x>"
    m.set_content("hello")
    m.add_alternative('<p>hello</p><img src="cid:logo@x">', subtype="html")
    m.get_payload()[1].add_related(
        LOGO, "image", "png", cid="<logo@x>", filename="logo.png", disposition="inline"
    )
    remote = RemoteMessage(
        uid=1, raw=m.as_bytes(), flags=frozenset(), internal_date=datetime.now(UTC)
    )
    result = await persist_message(
        db,
        user_id=user.id,
        account_id=account.id,
        mailbox_id=inbox.id,
        mailbox_special_use="inbox",
        remote=remote,
    )
    await db.commit()
    assert result is not None
    try:
        yield db, user.id, result.message.id
    finally:
        await db.execute(delete(User).where(User.id == user.id))
        await db.commit()
        await db.close()
        await engine.dispose()
        shutil.rmtree(Path(get_settings().blob_root) / str(user.id), ignore_errors=True)


async def test_quoted_images_travel_inside_the_reply(ctx) -> None:  # type: ignore[no-untyped-def]
    db, user_id, original = ctx
    gone = uuid.uuid4()
    html = (
        "<p>Thanks!</p><blockquote>"
        f'<img src="/api/messages/{original}/cid/logo@x" alt="logo">'
        f'<img src="http:///api/messages/{original}/cid/logo@x">'  # same image, mangled
        f'<img src="/api/messages/{gone}/cid/old@x">'  # its message is gone
        '<img src="https://example.com/banner.png"></blockquote>'
    )
    out, parts = await inline_images.embed(db, user_id, html, "example.com")
    assert len(parts) == 1 and parts[0].data == LOGO
    assert out.count(f'src="cid:{parts[0].cid}"') == 2
    assert "/api/messages/" not in out and 'alt="logo"' in out
    assert "https://example.com/banner.png" in out  # other images untouched

    account = SimpleNamespace(email="me@example.com", display_name="Me")
    draft = SimpleNamespace(
        payload={"to": ["ben@example.org"], "subject": "Re: Signed", "html": html}
    )
    raw, _, _, _ = build_mime(account, draft, None, [], html=out, inline=parts)  # type: ignore[arg-type]
    sent = email.message_from_bytes(raw, policy=policy.default)
    related = [p for p in sent.walk() if p.get_content_type() == "multipart/related"]
    assert related, "the HTML and its image belong together"
    image = next(p for p in sent.walk() if p.get_content_type() == "image/png")
    assert image["Content-ID"] == f"<{parts[0].cid}>"
    assert image.get_payload(decode=True) == LOGO


def test_links_already_sent_load_from_here() -> None:
    mid = uuid.uuid4()
    for src in (
        f"http:///api/messages/{mid}/cid/logo@x",
        f"https://mail.example.com/api/messages/{mid}/cid/logo@x",
        f"/api/messages/{mid}/cid/logo@x",
    ):
        result = sanitize_html(f'<img src="{src}">', uuid.uuid4(), allow_remote_images=True)
        assert f'src="/api/messages/{mid}/cid/logo@x"' in result.html, src
        assert not result.has_remote_images