ArchiveBox/archivebox/tests/test_persona_admin.py

import pytest
from typing import cast

from django.contrib.auth import get_user_model
from django.contrib.auth.models import UserManager
from django.urls import reverse

from archivebox.personas.importers import (
    PersonaImportResult,
    discover_persona_template_profiles,
    import_persona_from_source,
    resolve_browser_profile_source,
    resolve_custom_import_source,
)


pytestmark = pytest.mark.django_db

User = get_user_model()
ADMIN_HOST = "admin.archivebox.localhost:8000"


@pytest.fixture
def admin_user(db):
    return cast(UserManager, User.objects).create_superuser(
        username="personaadmin",
        email="personaadmin@test.com",
        password="testpassword",
    )


def _make_profile_source(tmp_path):
    user_data_dir = tmp_path / "Chrome User Data"
    profile_dir = user_data_dir / "Default"
    profile_dir.mkdir(parents=True)
    (profile_dir / "Preferences").write_text("{}")
    return resolve_browser_profile_source(
        browser="chrome",
        user_data_dir=user_data_dir,
        profile_dir="Default",
        browser_binary="/Applications/Google Chrome.app/Contents/MacOS/Google Chrome",
    )


def test_resolve_custom_import_source_accepts_exact_profile_dir(tmp_path):
    user_data_dir = tmp_path / "Brave User Data"
    profile_dir = user_data_dir / "Profile 2"
    profile_dir.mkdir(parents=True)
    (profile_dir / "Preferences").write_text("{}")

    source = resolve_custom_import_source(str(profile_dir))

    assert source.kind == "browser-profile"
    assert source.user_data_dir == user_data_dir.resolve()
    assert source.profile_dir == "Profile 2"


def test_resolve_custom_import_source_accepts_cdp_url():
    source = resolve_custom_import_source("ws://127.0.0.1:9222/devtools/browser/test-session")

    assert source.kind == "cdp"
    assert source.cdp_url == "ws://127.0.0.1:9222/devtools/browser/test-session"


def test_discover_persona_template_profiles_finds_chrome_profile_dirs(tmp_path):
    personas_dir = tmp_path / "personas"
    chrome_profile = personas_dir / "ExistingPersona" / "chrome_profile"
    default_profile = chrome_profile / "Default"
    default_profile.mkdir(parents=True)
    (default_profile / "Preferences").write_text("{}")

    discovered = discover_persona_template_profiles(personas_dir=personas_dir)

    assert len(discovered) == 1
    assert discovered[0].browser == "persona"
    assert discovered[0].source_name == "ExistingPersona"
    assert discovered[0].profile_dir == "Default"
    assert discovered[0].user_data_dir == chrome_profile.resolve()


def test_discover_persona_template_profiles_finds_home_abx_personas(monkeypatch, tmp_path):
    from archivebox.config.constants import CONSTANTS

    monkeypatch.setattr(CONSTANTS, "PERSONAS_DIR", tmp_path / "missing-data-personas")
    monkeypatch.setattr("archivebox.personas.importers.Path.home", lambda: tmp_path)

    chrome_profile = tmp_path / ".config" / "abx" / "personas" / "HomePersona" / "chrome_profile"
    default_profile = chrome_profile / "Default"
    default_profile.mkdir(parents=True)
    (default_profile / "Preferences").write_text("{}")

    discovered = discover_persona_template_profiles()

    assert len(discovered) == 1
    assert discovered[0].browser == "persona"
    assert discovered[0].source_name == "HomePersona"
    assert discovered[0].profile_dir == "Default"
    assert discovered[0].user_data_dir == chrome_profile.resolve()


def test_persona_admin_add_view_renders_import_ui(client, admin_user, monkeypatch, tmp_path):
    source = _make_profile_source(tmp_path)
    monkeypatch.setattr("archivebox.personas.forms.discover_local_browser_profiles", lambda: [source])
    monkeypatch.setattr("archivebox.personas.admin.discover_local_browser_profiles", lambda: [source])

    client.login(username="personaadmin", password="testpassword")
    response = client.get(reverse("admin:personas_persona_add"), HTTP_HOST=ADMIN_HOST)

    assert response.status_code == 200
    assert b"Bootstrap a persona from a real browser session" in response.content
    assert b"Google Chrome / Default" in response.content
    assert b"auth.json" in response.content


def test_import_persona_from_source_copies_user_agent_to_persona_config(admin_user, monkeypatch, tmp_path):
    from archivebox.personas.models import Persona

    source = _make_profile_source(tmp_path)
    persona = Persona.objects.create(name="AgentPersona", created_by=admin_user)

    def fake_export_browser_state(**kwargs):
        return True, {"user_agent": "Mozilla/5.0 Test Imported UA"}, "ok"

    monkeypatch.setattr("archivebox.personas.importers.export_browser_state", fake_export_browser_state)

    result = import_persona_from_source(
        persona,
        source,
        copy_profile=False,
        import_cookies=False,
        capture_storage=False,
    )

    persona.refresh_from_db()
    assert result.user_agent_imported is True
    assert persona.config["USER_AGENT"] == "Mozilla/5.0 Test Imported UA"


def test_persona_admin_add_post_runs_shared_importer(client, admin_user, monkeypatch, tmp_path):
    from archivebox.personas.models import Persona

    source = _make_profile_source(tmp_path)
    monkeypatch.setattr("archivebox.personas.forms.discover_local_browser_profiles", lambda: [source])
    monkeypatch.setattr("archivebox.personas.admin.discover_local_browser_profiles", lambda: [source])

    calls = {}

    def fake_import(persona, selected_source, **kwargs):
        calls["persona_name"] = persona.name
        calls["source"] = selected_source
        calls["kwargs"] = kwargs
        (persona.path / "cookies.txt").parent.mkdir(parents=True, exist_ok=True)
        (persona.path / "cookies.txt").write_text("# Netscape HTTP Cookie File\n")
        (persona.path / "auth.json").write_text('{"TYPE":"auth","cookies":[],"localStorage":{},"sessionStorage":{}}\n')
        return PersonaImportResult(
            source=selected_source,
            profile_copied=True,
            cookies_imported=True,
            storage_captured=True,
        )

    monkeypatch.setattr("archivebox.personas.forms.import_persona_from_source", fake_import)

    client.login(username="personaadmin", password="testpassword")
    response = client.post(
        reverse("admin:personas_persona_add"),
        {
            "name": "ImportedPersona",
            "created_by": str(admin_user.pk),
            "config": "{}",
            "import_mode": "discovered",
            "import_discovered_profile": source.choice_value,
            "import_copy_profile": "on",
            "import_extract_cookies": "on",
            "import_capture_storage": "on",
            "_save": "Save",
        },
        HTTP_HOST=ADMIN_HOST,
    )

    assert response.status_code == 302
    persona = Persona.objects.get(name="ImportedPersona")
    assert calls["persona_name"] == "ImportedPersona"
    assert calls["source"].profile_dir == "Default"
    assert calls["kwargs"] == {
        "copy_profile": True,
        "import_cookies": True,
        "capture_storage": True,
    }
    assert persona.COOKIES_FILE.endswith("cookies.txt")
    assert persona.AUTH_STORAGE_FILE.endswith("auth.json")