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


async def test_images_written_into_the_html_go_as_inline_parts() -> None:
    """A signature logo pasted into the editor is a data: URL, which Gmail does
    not show in mail it receives: it is sent as an inline part instead, once
    however often it appears."""
    import base64

    data_url = "data:image/png;base64," + base64.b64encode(LOGO).decode()
    html = (
        f'<p>Hi</p><img src="{data_url}" width="40"><p>--</p><img src="{data_url}">'
        '<img src="data:text/html;base64,PHNjcmlwdD4=">'  # not an image: left alone
    )
    out, parts = await inline_images.embed(None, uuid.uuid4(), html, "example.com")  # type: ignore[arg-type]
    assert len(parts) == 1
    assert parts[0].data == LOGO and parts[0].mime_type == "image/png"
    assert parts[0].filename == "image.png"
    assert out.count(f'src="cid:{parts[0].cid}"') == 2
    assert "data:image/png" not in out
    assert 'width="40"' in out

    account = SimpleNamespace(email="me@example.com", display_name="Me")
    draft = SimpleNamespace(payload={"to": ["a@example.org"], "subject": "Sig", "html": html})
    raw, _, _, _ = build_mime(account, draft, None, [], html=out, inline=parts)  # type: ignore[arg-type]
    msg = email.message_from_bytes(raw, policy=policy.default)
    images = [p for p in msg.walk() if p.get_content_type() == "image/png"]
    assert len(images) == 1 and images[0].get_content() == LOGO
    assert images[0]["Content-ID"] == f"<{parts[0].cid}>"


def test_smtp_servers_are_greeted_with_our_domain(monkeypatch) -> None:  # type: ignore[no-untyped-def]
    from app.config import smtp_helo_name

    settings = get_settings()
    monkeypatch.setattr(settings, "smtp_helo_name", "")
    assert smtp_helo_name() == settings.app_domain
    monkeypatch.setattr(settings, "smtp_helo_name", "mail.example.org")
    assert smtp_helo_name() == "mail.example.org"