re-xxm/tests/test_auth.py

# -*- coding: utf-8 -*-
"""auth 相关纯函数自测（不联网）: python tests/test_auth.py"""
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.stdout.reconfigure(encoding="utf-8")

import favsort as f  # noqa: E402

OLD = (
    "buvid3=OLD-B3; buvid4=OLD-B4; _uuid=OLD-UUID; buvid_fp=OLD-FP; "
    "b_nut=1700000000; SESSDATA=old%2C1700000000%2Cabc; bili_jct=oldcsrf; "
    "DedeUserID=28802351; DedeUserID__ckMd5=oldmd5; sid=oldsid; b_lsid=OLDL; "
    "bili_ticket=ticket; home_feed_column=5"
)


def test_parse_and_merge():
    pairs = f.parse_cookie(OLD)
    assert pairs[0] == ("buvid3", "OLD-B3")
    assert f.cookie_value(OLD, "SESSDATA") == "old%2C1700000000%2Cabc"

    new = {
        "SESSDATA": "new%2C1800000000%2Cxyz",
        "bili_jct": "newcsrf",
        "DedeUserID": "28802351",
        "DedeUserID__ckMd5": "newmd5",
        "sid": "newsid",
    }
    merged = f.merge_cookie(OLD, new)
    # 身份字段被替换
    assert f.cookie_value(merged, "SESSDATA") == "new%2C1800000000%2Cxyz"
    assert f.cookie_value(merged, "bili_jct") == "newcsrf"
    assert f.cookie_value(merged, "sid") == "newsid"
    # 设备指纹原样保留
    assert f.cookie_value(merged, "buvid3") == "OLD-B3"
    assert f.cookie_value(merged, "buvid4") == "OLD-B4"
    assert f.cookie_value(merged, "_uuid") == "OLD-UUID"
    assert f.cookie_value(merged, "buvid_fp") == "OLD-FP"
    assert f.cookie_value(merged, "b_nut") == "1700000000"
    assert f.cookie_value(merged, "home_feed_column") == "5"
    # 旧会话字段在没有新值时应被丢弃
    assert f.cookie_value(merged, "b_lsid") == ""


def test_set_cookie_and_url_fallback():
    headers = [
        "SESSDATA=abc%2C1800000000%2Cdef; Path=/; Domain=.bilibili.com; HttpOnly",
        "bili_jct=csrf2; Domain=.bilibili.com",
        "buvid3=SHOULD-NOT-OVERWRITE; Domain=.bilibili.com",
    ]
    got = f.cookies_from_set_cookie(headers)
    assert got["SESSDATA"] == "abc%2C1800000000%2Cdef"
    assert got["bili_jct"] == "csrf2"

    merged = f.merge_cookie(OLD, got)
    assert f.cookie_value(merged, "buvid3") == "OLD-B3"   # 设备字段不被 Set-Cookie 覆盖


def test_expiry():
    assert f.sessdata_expiry(OLD) == 1700000000
    assert f.sessdata_expiry("SESSDATA=bad") is None
    assert f.sessdata_expiry("") is None


def test_ensure_device_cookies():
    called = {"n": 0}

    def fake_get_json(cfg, url, auth=False):
        called["n"] += 1
        return {"code": 0, "data": {"b_3": "NEW-B3", "b_4": "NEW-B4"}}

    orig = f.get_json
    f.get_json = fake_get_json
    try:
        cfg = {"cookie": "SESSDATA=x; bili_jct=y"}
        added = f.ensure_device_cookies(cfg)
        assert "buvid3" in added and "buvid4" in added and "_uuid" in added
        assert f.cookie_value(cfg["cookie"], "buvid3") == "NEW-B3"
        assert f.cookie_value(cfg["cookie"], "buvid4") == "NEW-B4"
        assert f.cookie_value(cfg["cookie"], "_uuid").endswith("infoc")

        # 已有设备字段时不应调用 spi，也不应改动
        called["n"] = 0
        cfg2 = {"cookie": OLD}
        added2 = f.ensure_device_cookies(cfg2)
        assert added2 == [] and called["n"] == 0
        assert cfg2["cookie"] == OLD
    finally:
        f.get_json = orig


def test_finish_auth_write_and_rollback():
    import json
    import tempfile

    new = {"SESSDATA": "new%2C1800000000%2Cxyz", "bili_jct": "newcsrf",
           "DedeUserID": "28802351", "DedeUserID__ckMd5": "newmd5"}
    cfg = {"mid": 28802351, "csrf": "oldcsrf", "cookie": OLD,
           "overrides": {}, "sleep": {}}
    old_raw = json.dumps(cfg, ensure_ascii=False, indent=2)

    orig = (f.CONFIG_PATH, f.CONFIG_BACKUP, f.fetch_nav, f.ensure_device_cookies)

    def fake_nav_ok(c):
        return {"code": 0, "data": {"isLogin": True, "uname": "tester", "mid": 28802351}}

    def fake_nav_bad(c):
        return {"code": -101, "message": "账号未登录", "data": {"isLogin": False}}

    def fake_device(c):
        return []

    try:
        with tempfile.TemporaryDirectory() as d:
            path = os.path.join(d, "config.json")
            f.CONFIG_PATH, f.CONFIG_BACKUP = path, path + ".bak"
            f.ensure_device_cookies = fake_device

            with open(path, "w", encoding="utf-8") as fh:
                fh.write(old_raw)

            # 成功：写入新身份字段，设备字段不变，并生成备份
            f.fetch_nav = fake_nav_ok
            assert f.finish_auth(cfg, dict(new)) == 0
            with open(path, encoding="utf-8") as fh:
                saved = json.load(fh)
            assert "newcsrf" == saved["csrf"]
            assert f.cookie_value(saved["cookie"], "buvid3") == "OLD-B3"
            assert f.cookie_value(saved["cookie"], "SESSDATA") == "new%2C1800000000%2Cxyz"
            assert os.path.exists(path + ".bak")

            # 失败：用上一份备份回滚
            before = open(path, encoding="utf-8").read()
            cfg2 = dict(saved)
            f.fetch_nav = fake_nav_bad
            assert f.finish_auth(cfg2, {"SESSDATA": "x%2C1%2Cy", "bili_jct": "z"}) == 4
            after = open(path, encoding="utf-8").read()
            assert after == before
    finally:
        f.CONFIG_PATH, f.CONFIG_BACKUP, f.fetch_nav, f.ensure_device_cookies = orig


def test_qr_png():
    import tempfile

    path = os.path.join(tempfile.gettempdir(), "favsort-test-qr.png")
    try:
        f.write_qr_png("https://example.com/login?key=abc&from=test", path)
        with open(path, "rb") as fh:
            data = fh.read()
        assert data.startswith(b"\x89PNG\r\n\x1a\n")
        assert len(data) > 500
    finally:
        if os.path.exists(path):
            os.remove(path)


def main():
    test_parse_and_merge()
    test_set_cookie_and_url_fallback()
    test_expiry()
    test_ensure_device_cookies()
    test_finish_auth_write_and_rollback()
    test_qr_png()
    print("OK: auth 纯函数测试全部通过")


if __name__ == "__main__":
    main()