live-hub-py/tests/test_huya_app_login.py

"""虎牙 App 密码登录及相关组件测试。"""

import base64
import os
import struct
import unittest
from unittest.mock import patch, MagicMock

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

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from core.huya import (
    HuyaAppLoginError,
    HuyaAppPasswordLogin,
    HuyaAppQrAuthRequiredError,
    login_huya_app_password,
    login_huya_password,
)
from core.huya.cert_forge import build_p1, decrypt_cert, forge_cert, parse_p1
from core.huya.device_profile import generate_profile, get_profile
from core.huya.envelope_forge import Envelope
from core.huya.login import HuyaLoginResult
from core.huya.nonce_forge import K1_DEFAULT, gen_nonce
from core.huya.udb_aes import udb_decrypt, udb_encrypt
from core.huya.wup_encoder import build_password_login_wup

from web.backend.database import Base
from web.backend.models import User, HuyaAccount
from web.backend.schemas import HuyaAppPasswordLoginRequest, HuyaPasswordLoginSelectedRequest
from web.backend.routers.huya import (
    app_password_login_account,
    app_password_login_selected_accounts,
)


class TestHuyaAppLogin(unittest.TestCase):

    def setUp(self):
        self.engine = create_engine("sqlite://")
        Base.metadata.create_all(self.engine)
        self.session = sessionmaker(bind=self.engine)()
        self.admin = User(username="admin", password_hash="hash", role="super_admin")
        self.session.add(self.admin)
        self.session.commit()

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

    def test_udb_aes_roundtrip(self):
        key = b"4VYcPdvKKqjBHZtC"
        plain = b"Hello, Huya App Login Protocol!"
        enc = udb_encrypt(key, plain)
        dec = udb_decrypt(key, enc)
        self.assertEqual(dec, plain)

    def test_cert_and_nonce_forge(self):
        uid = 1199666914671
        rnd = gen_nonce(uid, K1_DEFAULT)
        self.assertEqual(len(rnd), 20)

        fp = b"02df398797432eadefcc12767119ad5e80999389"
        cred = b"\x0a" + os.urandom(113)
        p1 = build_p1(b"5008", fp, cred, rnd=rnd)
        self.assertEqual(len(p1), 187)

        cert = forge_cert(p1)
        dec = decrypt_cert(cert)
        parsed = parse_p1(dec)
        self.assertEqual(parsed["rnd"], rnd)
        self.assertEqual(parsed["fingerprint"], fp)
        self.assertEqual(parsed["cred"], cred)

    def test_envelope_patching(self):
        env = Envelope.load()
        self.assertGreater(env.uid, 0)
        self.assertEqual(len(env.cert_b64), 260)

        new_uid = 1199666911746
        env.patch_uid(new_uid)
        self.assertEqual(env.uid, new_uid)

        fake_cert = bytes([0x0C, 0x20]) + os.urandom(192)
        env.patch_cert(fake_cert)
        self.assertEqual(len(env.cert_b64), 260)

        wup_b64 = env.wup_b64()
        self.assertGreater(len(wup_b64), 0)

    def test_device_profile_generation(self):
        p1 = generate_profile()
        self.assertEqual(p1["os"], "android")
        self.assertEqual(len(p1["fingerprint"]), 40)
        self.assertEqual(len(p1["device_id"]), 40)
        self.assertEqual(p1["hdid"], "ed0db8334cadd236c00cadf7e11ab5a5")

        p2 = get_profile("test_user_account_123")
        p3 = get_profile("test_user_account_123")
        self.assertEqual(p2["fingerprint"], p3["fingerprint"])

    def test_wup_encoder_output(self):
        dev = generate_profile()
        pkt = build_password_login_wup(
            uid_str="300023887",
            sha1_password="772ed992b0e161276f44ec63671e60155c506294",
            safedeviceid=dev["safedeviceid"],
            hdid=dev["hdid"],
            session=3251699,
            trace_id="test-trace-id",
            user_action_json="{}",
            device_info=dev,
        )
        self.assertGreater(len(pkt), 500)
        total_len = struct.unpack(">I", pkt[:4])[0]
        self.assertEqual(total_len, len(pkt))

    def test_router_functions(self):
        mock_res = HuyaLoginResult(
            success=True,
            cookie="udb_cred=mock_cred_123; yyuid=1199666914671; udb_uid=mock_uid; username=mock_test_huya_user",
            message="App协议登录成功",
            sdid="mock_sdid_123",
        )

        with patch("web.backend.routers.huya.login_huya_app_password", return_value=mock_res):
            req = HuyaAppPasswordLoginRequest(
                username="mock_test_huya_user",
                password="mock_password_123",
                tag="test_tag",
            )
            resp = app_password_login_account(req=req, db=self.session, current=self.admin)
            self.assertTrue(resp["success"])
            self.assertIn("App", resp["message"])
            self.assertEqual(resp["account"].username, "mock_test_huya_user")

            # Test batch app login router function
            acc = self.session.query(HuyaAccount).filter(HuyaAccount.username == "mock_test_huya_user").first()
            batch_req = HuyaPasswordLoginSelectedRequest(account_ids=[acc.id])
            batch_resp = app_password_login_selected_accounts(req=batch_req, db=self.session, current=self.admin)
            self.assertTrue(batch_resp["success"])
            self.assertEqual(batch_resp["count"], 1)
            self.assertEqual(len(batch_resp["results"]), 1)
            self.assertTrue(batch_resp["results"][0]["success"])


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