live-hub-py/tests/test_cookie_operations.py

import os
import unittest
from types import SimpleNamespace
from unittest.mock import patch

os.environ.setdefault("DATABASE_URL", "sqlite://")
os.environ.setdefault("APP_ENCRYPTION_KEY", "MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3ODlhYmNkZWY=")

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from web.backend.database import Base
from web.backend.models import Account, LoginTask, User
from web.backend.routers.cookies import (
    check_cookie_operations,
    get_cookie,
    list_cookie_operation_tags,
    list_cookie_operations,
    list_cookies,
)


class CookieOperationTests(unittest.TestCase):
    def setUp(self):
        self.engine = create_engine("sqlite://")
        Base.metadata.create_all(self.engine)
        self.session = sessionmaker(bind=self.engine)()

        self.support = User(username="support", password_hash="hash", role="support")
        self.other_support = User(username="other", password_hash="hash", role="support")
        self.session.add_all([self.support, self.other_support])
        self.session.commit()

        owned = Account(
            username="owned-account",
            password="owned-password",
            email="owned@example.com",
            email_password="mail-password",
            assigned_to=self.support.id,
        )
        other = Account(
            username="other-account",
            password="other-password",
            email="other@example.com",
            email_password="mail-password",
            assigned_to=self.other_support.id,
        )
        self.session.add_all([owned, other])
        self.session.commit()
        self.session.add_all([
            LoginTask(batch_id="owned", account_id=owned.id, created_by=self.support.id, status="success", cookie="owned-secret"),
            LoginTask(batch_id="other", account_id=other.id, created_by=self.other_support.id, status="success", cookie="other-secret"),
        ])
        self.session.commit()
        self.owned_task = self.session.query(LoginTask).filter(LoginTask.batch_id == "owned").one()
        self.other_task = self.session.query(LoginTask).filter(LoginTask.batch_id == "other").one()

    def tearDown(self):
        self.session.close()
        Base.metadata.drop_all(self.engine)
        self.engine.dispose()

    def test_support_operation_list_is_scoped_and_never_returns_credentials(self):
        result = list_cookie_operations(
            search="",
            tag="",
            page=1,
            page_size=20,
            db=self.session,
            current=self.support,
        )

        self.assertEqual(result["total"], 1)
        self.assertEqual(result["items"], [{
            "id": self.owned_task.id,
            "account_username": "owned-account",
            "tag": "",
            "ck_check_status": "",
            "ck_checked_at": None,
            "created_at": None,
        }])
        self.assertNotIn("owned-secret", str(result))
        self.assertNotIn("owned-password", str(result))

    def test_support_cannot_use_cookie_management_endpoints(self):
        with self.assertRaisesRegex(Exception, "cookie:view"):
            list_cookies(db=self.session, current=self.support)
        with self.assertRaisesRegex(Exception, "cookie:view"):
            get_cookie(self.owned_task.id, db=self.session, current=self.support)

    def test_cookie_list_filters_by_account_tag(self):
        admin = User(username="admin", password_hash="hash", role="super_admin")
        self.session.add(admin)
        self.session.query(Account).filter(Account.id == self.owned_task.account_id).update({"tag": "A组"})
        self.session.query(Account).filter(Account.id == self.other_task.account_id).update({"tag": "B组"})
        self.session.commit()

        result = list_cookies(
            search="",
            tag="A组",
            account_names="",
            page=1,
            page_size=20,
            include_cookie=False,
            db=self.session,
            current=admin,
        )

        self.assertEqual(result["total"], 1)
        self.assertEqual([item["id"] for item in result["items"]], [self.owned_task.id])

    def test_support_operation_tag_filter_and_list_stay_scoped(self):
        self.session.query(Account).filter(Account.id == self.owned_task.account_id).update({"tag": "我的标签"})
        self.session.query(Account).filter(Account.id == self.other_task.account_id).update({"tag": "别人的标签"})
        self.session.commit()

        tags = list_cookie_operation_tags(db=self.session, current=self.support)
        result = list_cookie_operations(
            search="",
            tag="别人的标签",
            page=1,
            page_size=20,
            db=self.session,
            current=self.support,
        )

        self.assertEqual(tags, ["我的标签"])
        self.assertEqual(result["total"], 0)

    @patch("web.backend.routers.cookies._check_one_cookie")
    def test_support_can_check_only_assigned_cookie(self, check_one_cookie):
        check_one_cookie.return_value = {
            "id": self.owned_task.id,
            "valid": True,
            "message": "有效",
            "fish_ball": 9,
            "nickname": "tester",
            "level": 1,
            "checked_at": "2026-01-01T00:00:00+00:00",
        }

        result = check_cookie_operations(
            ids=f"{self.owned_task.id},{self.other_task.id}",
            db=self.session,
            current=self.support,
        )

        self.assertEqual([item["id"] for item in result["results"]], [self.owned_task.id])
        self.session.refresh(self.owned_task)
        self.session.refresh(self.other_task)
        self.assertEqual(self.owned_task.ck_check_status, "valid")
        self.assertEqual(self.other_task.ck_check_status, "")

    def test_operation_permission_does_not_allow_unrelated_users(self):
        no_permission_user = SimpleNamespace(
            id=999,
            username="no-permission",
            role="support",
            custom_permissions=[],
        )
        with self.assertRaisesRegex(Exception, "cookie:operate"):
            list_cookie_operations(
                search="",
                tag="",
                page=1,
                page_size=20,
                db=self.session,
                current=no_permission_user,
            )


if __name__ == "__main__":
    unittest.main()