live-hub-py/tests/test_cookie_operations.py

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

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.services.login_service import LoginBatchRunner, cleanup_orphan_relogin_tasks
from web.backend.services.cookie_check_service import check_douyu_cookie
from web.backend.routers.cookies import (
    check_cookie_operations,
    get_cookie,
    list_cookie_operation_tags,
    list_cookie_operations,
    list_cookies,
    relogin_invalid_cookie_operations,
)


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,
            "relogin_status": "",
            "relogin_message": "",
            "relogin_batch_id": "",
        }])
        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,
            )

    def test_relogin_failed_record_stays_visible_and_can_be_checked(self):
        self.owned_task.status = "relogin_failed"
        self.owned_task.message = "重新登录失败: 密码错误（旧 Cookie 已保留）"
        self.session.commit()

        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"][0]["relogin_status"], "relogin_failed")
        self.assertIn("旧 Cookie 已保留", result["items"][0]["relogin_message"])

    def test_service_restart_cleans_orphan_relogin_tasks(self):
        self.owned_task.status = "relogin_running"
        self.session.commit()

        cleaned = cleanup_orphan_relogin_tasks(self.session)

        self.assertEqual(cleaned, 1)
        self.session.refresh(self.owned_task)
        self.session.refresh(self.other_task)
        self.assertEqual(self.owned_task.status, "relogin_failed")
        self.assertIn("服务重启", self.owned_task.message)
        self.assertEqual(self.other_task.status, "success")

    def test_runner_copies_proxy_config_before_background_execution(self):
        proxy = SimpleNamespace(
            enabled=True,
            http="http://127.0.0.1:8080",
            https="",
            api_url="",
            whitelist_enabled=False,
            whitelist_platform=None,
            whitelist_credentials=None,
            whitelist_uid="",
            whitelist_ukey="",
        )
        runner = LoginBatchRunner(
            db=self.session,
            account_ids=[],
            created_by=self.support.id,
            creator_permissions=[],
            proxy_config=proxy,
        )
        proxy.http = "http://changed.example:8080"

        proxy_dict, _ = runner._resolve_static_proxy()

        self.assertEqual(proxy_dict, {
            "http": "http://127.0.0.1:8080",
            "https": "http://127.0.0.1:8080",
        })

    @patch("web.backend.services.cookie_check_service.requests.get")
    def test_cookie_check_uses_new_cookie_result(self, mock_get):
        mock_get.side_effect = [
            SimpleNamespace(json=lambda: {"error": 0, "data": {"count": 9}}),
            SimpleNamespace(json=lambda: {"error": 0, "data": {"nn": "new-name", "lv": 12}}),
        ]

        result = check_douyu_cookie("new-cookie")

        self.assertTrue(result["valid"])
        self.assertEqual(result["fish_ball"], 9)
        self.assertEqual(result["nickname"], "new-name")
        self.assertEqual(result["level"], 12)

    @patch("web.backend.routers.cookies._start_relogin_tasks")
    def test_relogin_invalid_only_targets_support_visible_accounts(self, start_relogin):
        self.owned_task.ck_check_status = "invalid"
        self.other_task.ck_check_status = "invalid"
        self.session.commit()
        start_relogin.return_value = {"batch_id": "batch", "count": 1, "skipped": 0, "success": True}

        result = relogin_invalid_cookie_operations(
            search="", tag="", db=self.session, current=self.support,
        )

        self.assertTrue(result["success"])
        selected_tasks = start_relogin.call_args.args[0]
        self.assertEqual([task.id for task in selected_tasks], [self.owned_task.id])


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