ksjsb/tests/test_main.py
2026-07-30 20:25:56 +08:00

557 lines
21 KiB
Python

import base64
import json
import os
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch
from urllib.parse import parse_qs, urlsplit
from core.captured_profile import CapturedProfile
from core.sig import sig
from main import (
BuiltRequest,
HarTemplateStore,
HttpResult,
KsTaskRunner,
SignedClient,
apply_cookie_overrides_to_pairs,
build_parser,
decode_base64_json,
extract_ad_material,
extract_task_items,
load_account_from_args,
missing_ad_material_fields,
parse_account_spec,
parse_reward_template_text,
print_normal_output,
profile_from_account,
require_account_from_args,
resolve_har_files,
replace_pair_values,
)
class FakeSig3State:
def __init__(self):
self.inputs = []
def sig3_hex(self, input_value, **_kwargs):
self.inputs.append(input_value)
return "S3"
class MainClientTests(unittest.TestCase):
def test_replace_pair_values_preserves_order_and_appends_missing(self):
pairs = [("a", "1"), ("b", "2"), ("a", "3")]
self.assertEqual(
replace_pair_values(pairs, {"a": "9", "c": "4"}),
[("a", "9"), ("b", "2"), ("a", "9"), ("c", "4")],
)
def test_decode_base64_json_accepts_urlsafe_unpadded_input(self):
payload = {"pageId": 1, "neoParams": "abc"}
raw = json.dumps(payload, separators=(",", ":")).encode()
encoded = base64.urlsafe_b64encode(raw).decode().rstrip("=")
self.assertEqual(decode_base64_json(encoded), payload)
def test_build_form_request_adds_all_native_signatures(self):
profile = CapturedProfile(
client_salt="salt",
api_st="token",
did="DID",
odid="ODID",
rdid="RDID",
egid="EGID",
client_key="ck",
)
state = FakeSig3State()
client = SignedClient(
profile=profile,
sig3_state=state,
xfalcon_func=lambda _data: "XF",
tokensig_func=lambda sig_value, _salt: "TK_" + sig_value[:6],
)
request = client.build_form_request(
"api.e.kuaishou.com",
"/rest/test",
[("foo", "bar")],
query_pairs=[("z", "2"), ("sig", "old"), ("__NS_sig3", "old")],
include_common_body=False,
)
query = parse_qs(urlsplit(request.url).query)
body = parse_qs(request.body_text)
expected_sig = sig({"z": "2", "foo": "bar"})
self.assertEqual(body["foo"], ["bar"])
self.assertEqual(query["sig"], [expected_sig])
self.assertEqual(query["__NS_sig3"], ["S3"])
self.assertEqual(query["__NStokensig"], ["TK_" + expected_sig[:6]])
self.assertEqual(query["__NS_xfalcon"], ["XF"])
self.assertEqual(state.inputs, ["/rest/test" + expected_sig])
def test_parse_reward_template_text_extracts_latest_reward_material(self):
str_e = b"reward-input"
body = "client_key=ck&encData=old&sign=" + ("a" * 64)
sig3_value = "41" * 24
log_text = "\n".join(
[
"[STR][atlasEncrypt in.raw] inputB64="
+ base64.b64encode(str_e).decode(),
"[STR][atlasSign out] len=64 ret=" + ("b" * 64),
"[REQ][reward] url=https://api.e.kuaishou.com/rest/e/reward/mixed/ad"
+ "?a=1&__NS_sig3="
+ sig3_value
+ " bodyB64="
+ base64.b64encode(body.encode()).decode(),
]
)
template = parse_reward_template_text(log_text)
self.assertEqual(template.str_e_bytes, str_e)
self.assertEqual(template.body_sign, "b" * 64)
self.assertEqual(template.query_dict["a"], "1")
self.assertEqual(template.query_dict["__NS_sig3"], sig3_value)
self.assertEqual(template.body_dict["client_key"], "ck")
def test_extract_task_items_walks_common_task_lists(self):
payload = {
"data": {
"dailyTasks": [{"name": "watch ad"}],
"taskCards": [{"taskName": "card", "subTasks": [{"title": "sub"}]}],
"list": [{"notTask": True}],
}
}
names = [item.get("name") or item.get("taskName") or item.get("title") for item in extract_task_items(payload)]
self.assertEqual(names, ["watch ad", "card", "sub"])
def test_extract_task_items_ignores_generic_non_task_lists(self):
payload = {
"data": {
"dailyTasks": [{"name": "task"}],
"someWidget": {"list": [{"url": "x"}, {"text": "not a task"}]},
}
}
self.assertEqual(len(extract_task_items(payload)), 1)
def test_extract_ad_material_accepts_root_level_reward_response(self):
neo_payload = {"businessId": 672, "extParams": "EXT", "pageId": 11101, "subPageId": 100026367}
neo_params = base64.b64encode(json.dumps(neo_payload, separators=(",", ":")).encode()).decode()
payload = {
"llsid": 123,
"feeds": [
{
"liveStreamId": "feed-1",
"ad": {
"llsid": 123,
"creativeId": 456,
"posId": "24067",
"adDataV2": {"inspireAdInfo": {"neoParams": neo_params, "adExtInfo": "AD"}},
},
}
],
}
material = extract_ad_material(payload)
self.assertEqual(material.neo_params, neo_params)
self.assertEqual(material.neo_ext_params, "EXT")
self.assertEqual(material.llsid, 123)
self.assertEqual(material.creative_id, 456)
self.assertEqual(material.feed_id, "feed-1")
self.assertEqual(material.ad_ext_info, "AD")
def test_har_template_store_extracts_query_and_headers(self):
har = {
"log": {
"entries": [
{
"request": {
"method": "GET",
"url": (
"https://nebula.kuaishou.com/rest/n/nebula/activity/earn/overview/tasks"
"?source=bottom_guide_first&bizParams=%7B%22refreshType%22%3A%22init%22%7D"
),
"headers": [
{"name": "Host", "value": "nebula.kuaishou.com"},
{"name": "Content-Length", "value": "123"},
{"name": "Referer", "value": "https://nebula.kuaishou.com/nebula/earning"},
{"name": "Cookie", "value": "kuaishou.api_st=old"},
{"name": "Accept-Encoding", "value": "gzip, deflate, br"},
],
}
}
]
}
}
with tempfile.TemporaryDirectory() as tmp:
path = Path(tmp) / "sample.har"
path.write_text(json.dumps(har), encoding="utf-8")
store = HarTemplateStore.from_har_files([path])
template = store.get("nebula.kuaishou.com", "/rest/n/nebula/activity/earn/overview/tasks")
self.assertIsNotNone(template)
self.assertEqual(template.query_pairs[0], ("source", "bottom_guide_first"))
self.assertEqual(template.query_pairs[1], ("bizParams", '{"refreshType":"init"}'))
headers = template.sanitized_headers(use_cookie=False)
self.assertEqual(headers["Host"], "nebula.kuaishou.com")
self.assertEqual(headers["Referer"], "https://nebula.kuaishou.com/nebula/earning")
self.assertEqual(headers["Accept-Encoding"], "gzip")
self.assertNotIn("Content-Length", headers)
self.assertNotIn("Cookie", headers)
def test_plain_get_uses_har_template_headers_without_stale_cookie(self):
store = HarTemplateStore()
store.put(
"nebula.kuaishou.com",
"/rest/example",
"GET",
[("source", "har")],
[
("Referer", "https://nebula.kuaishou.com/nebula/earning"),
("Cookie", "old=1"),
("Accept-Encoding", "gzip, deflate, br"),
],
)
client = SignedClient(template_store=store)
request = client.build_plain_get_request(
"nebula.kuaishou.com",
"/rest/example",
client.template_query_pairs("nebula.kuaishou.com", "/rest/example", []),
)
self.assertIn("source=har", request.url)
self.assertEqual(request.headers["Referer"], "https://nebula.kuaishou.com/nebula/earning")
self.assertEqual(request.headers["Host"], "nebula.kuaishou.com")
self.assertEqual(request.headers["Accept-Encoding"], "gzip")
self.assertNotIn("Cookie", request.headers)
def test_execute_returns_structured_http_error(self):
class FailingClient:
def send(self, request):
return HttpResult(False, 500, request.url, "boom", None, request)
request = BuiltRequest("GET", "https://example.test/path", {}, None)
runner = KsTaskRunner(FailingClient(), dry_run=False)
result = runner._execute("probe", request)
self.assertFalse(result["ok"])
self.assertEqual(result["event"], "probe")
self.assertEqual(result["status"], 500)
self.assertEqual(result["error"], "http_status_500")
self.assertEqual(result["responseText"], "boom")
def test_watch_ad_skips_followups_when_real_ad_material_is_incomplete(self):
class ProbeRunner(KsTaskRunner):
def __init__(self):
super().__init__(client=object(), dry_run=False)
self.called = []
def fetch_reward_ad(self):
return {
"ok": True,
"event": "reward-mixed-ad",
"responseJson": {"result": 1, "feeds": []},
}
def neo_task_info(self, neo_params):
self.called.append("neo_task_info")
return {}
def report_ad_task(self, material):
self.called.append("report_ad_task")
return {}
def nebula_reward(self, material):
self.called.append("nebula_reward")
return {}
def activity_reward(self, material):
self.called.append("activity_reward")
return {}
def coin_send(self, material):
self.called.append("coin_send")
return {}
runner = ProbeRunner()
result = runner.watch_ad_once()
self.assertEqual(runner.called, [])
self.assertIn("missing_ad_material", result["warnings"])
def test_ad_material_accepts_reward_ad_without_neo_params(self):
payload = {
"result": 1,
"llsid": 2006042943124939954,
"feeds": [
{
"photo_id": "5194902358263689735",
"ad": {
"llsid": 2006042943124939954,
"creativeId": 211282409779,
"pageId": 11101,
"subPageId": 100026367,
"displayType": 2,
"adDataV2": {
"posId": 24067,
"inspireAdInfo": {"adExtInfo": "AD_EXT"},
},
"extData": json.dumps({"awardCoin": 46}),
},
}
],
}
material = extract_ad_material(payload, allow_fallback=False)
self.assertEqual(material.llsid, 2006042943124939954)
self.assertEqual(material.creative_id, 211282409779)
self.assertEqual(material.feed_id, "5194902358263689735")
self.assertEqual(material.display_type, 2)
self.assertEqual(material.ad_ext_info, "AD_EXT")
self.assertEqual(missing_ad_material_fields(material), [])
def test_sign_in_skips_report_when_resource_has_no_event(self):
class ProbeRunner(KsTaskRunner):
def __init__(self):
super().__init__(client=object(), dry_run=False)
self.reported = False
def get_sign_in_resource(self):
return {"ok": True, "responseJson": {"result": 1, "data": {"signInUnionSpecialAreaData": None}}}
def _execute(self, event, request):
if event == "sign-in-report":
self.reported = True
return {}
runner = ProbeRunner()
result = runner.sign_in()
self.assertFalse(runner.reported)
self.assertTrue(result["skipped"])
self.assertEqual(result["reason"], "no_sign_in_event")
def test_treasure_box_skips_report_when_info_is_unavailable(self):
class ProbeRunner(KsTaskRunner):
def __init__(self):
super().__init__(client=object(), dry_run=False)
self.reported = False
def treasure_box_info(self):
return {"ok": True, "responseJson": {"result": 102512, "msg": ""}}
def _execute(self, event, request):
if event == "treasure-box-report":
self.reported = True
return {}
runner = ProbeRunner()
result = runner.open_treasure_box()
self.assertFalse(runner.reported)
self.assertTrue(result["skipped"])
self.assertEqual(result["reason"], "treasure_box_unavailable")
def test_treasure_box_report_uses_h5_json_endpoint(self):
class FakeClient:
def __init__(self):
self.request = None
def build_plain_json_request(self, host, path, payload, query_pairs=None):
self.request = BuiltRequest(
"POST",
f"https://{host}{path}",
{"Host": host, "Content-Type": "application/json"},
json.dumps(payload).encode(),
{"query": query_pairs or []},
)
return self.request
def template_query_pairs(self, host, path, fallback):
return fallback
class ProbeRunner(KsTaskRunner):
def treasure_box_info(self):
return {
"ok": True,
"responseJson": {
"result": 1,
"data": {
"eventTrackingLogInfo": {
"eventTrackingTaskId": 20035,
"resourceId": "externalFeed_externalTreasureChestWidget",
}
},
},
}
def _execute(self, event, request):
return {"event": event, "request": request.as_dict(redact=False)}
client = FakeClient()
runner = ProbeRunner(client=client, dry_run=False)
result = runner.open_treasure_box()
self.assertEqual(result["event"], "treasure-box-report")
self.assertEqual(client.request.url, "https://nebula.kuaishou.com/rest/wd/encourage/unionTask/treasureBox/report")
self.assertIn("eventTrackingLogInfo", client.request.body_text)
def test_watch_ad_reports_business_error_without_missing_material(self):
class ProbeRunner(KsTaskRunner):
def __init__(self):
super().__init__(client=object(), dry_run=False)
def fetch_reward_ad(self):
return {
"ok": True,
"event": "reward-mixed-ad",
"responseJson": {"result": 6001, "message": "no ad"},
}
result = ProbeRunner().watch_ad_once()
self.assertIn("reward_business_error_6001", result["warnings"])
self.assertNotIn("missing_ad_material", result["warnings"])
def test_reward_ad_uses_account_overridden_query_not_old_capture_query(self):
body = "client_key=ck&encData=old&sign=" + ("a" * 64)
log_text = "\n".join(
[
"[STR][atlasEncrypt in.raw] inputB64=" + base64.b64encode(b"reward-input").decode(),
"[STR][atlasSign out] len=64 ret=" + ("b" * 64),
"[REQ][reward] url=https://api.e.kuaishou.com/rest/e/reward/mixed/ad"
+ "?ud=OLD_USER&did=DID_OLD&egid=EGID_OLD"
+ " bodyB64="
+ base64.b64encode(body.encode()).decode(),
]
)
template = parse_reward_template_text(log_text)
class FakeRewardSign:
def reward_sign(self, *_args, **_kwargs):
return "SIGN_NEW"
class FakeClient:
def __init__(self):
self.base_query_pairs = [("ud", "NEW_USER"), ("did", "DID_NEW"), ("egid", "EGID_NEW")]
self.sig3_state = FakeRewardSign()
self.query_pairs = None
def load_kwsg_tables(self):
return b"t1", b"t2"
def build_form_request(self, _host, _path, _body_pairs, *, query_pairs=None, include_common_body=True):
self.query_pairs = query_pairs
return BuiltRequest("POST", "https://api.e.kuaishou.com/rest/e/reward/mixed/ad", {}, b"")
client = FakeClient()
runner = KsTaskRunner(client=client, reward_template=template, dry_run=True, redact=False)
with patch("main.kwsg_10400_raw", return_value=b"ENC"):
runner.fetch_reward_ad()
self.assertEqual(client.query_pairs, client.base_query_pairs)
def test_account_spec_overrides_profile_and_request_pairs(self):
account = parse_account_spec(
"大号#userId=1579452490; ud=1579452490; did=ANDROID_new; "
"kuaishou.api_st=TOKEN; token=TOKEN; egid=EGID_new; "
"oDid=ODID_new; rdid=RDID_new; client_key=client_new; appver=14.0#salt_new"
)
profile = profile_from_account(account)
self.assertEqual(profile.client_salt, "salt_new")
self.assertEqual(profile.api_st, "TOKEN")
self.assertEqual(profile.did, "ANDROID_new")
self.assertEqual(profile.odid, "ODID_new")
self.assertEqual(profile.rdid, "RDID_new")
self.assertEqual(profile.egid, "EGID_new")
self.assertEqual(profile.client_key, "client_new")
pairs = apply_cookie_overrides_to_pairs(
[("ud", "old"), ("did", "old"), ("appver", "old"), ("kept", "1")],
account.cookie_dict,
)
self.assertEqual(
pairs,
[("ud", "1579452490"), ("did", "ANDROID_new"), ("appver", "14.0"), ("kept", "1")],
)
def test_parser_defaults_to_run_all_without_cli_noise(self):
with patch.dict(os.environ, {}, clear=True):
args = build_parser().parse_args([])
self.assertEqual(args.command, "run-all")
def test_parser_help_hides_debug_options(self):
with patch.dict(os.environ, {}, clear=True):
help_text = build_parser().format_help()
self.assertIn("run-all", help_text)
self.assertIn("coin", help_text)
self.assertNotIn("--har", help_text)
self.assertNotIn("--reward-log", help_text)
self.assertNotIn("--neo-params", help_text)
self.assertNotIn("task-center", help_text)
def test_env_account_is_used_when_cli_account_is_empty(self):
with patch.dict(
os.environ,
{
"KS_ACCOUNT": "大号#userId=1; did=DID_ENV; kuaishou.api_st=TOKEN_ENV#salt_env",
},
clear=True,
):
args = build_parser().parse_args([])
account = load_account_from_args(args)
self.assertIsNotNone(account)
self.assertEqual(account.remark, "大号")
self.assertEqual(account.cookie_dict["did"], "DID_ENV")
self.assertEqual(account.client_salt, "salt_env")
def test_normal_cli_requires_account_config(self):
with patch.dict(os.environ, {}, clear=True):
args = build_parser().parse_args([])
with self.assertRaises(SystemExit) as caught:
require_account_from_args(args)
self.assertIn("KS_ACCOUNT", str(caught.exception))
def test_har_is_not_loaded_by_default(self):
with patch.dict(os.environ, {}, clear=True):
args = build_parser().parse_args([])
self.assertEqual(resolve_har_files(args), [])
with patch.dict(os.environ, {"KS_USE_HAR": "1"}, clear=True):
args = build_parser().parse_args([])
self.assertTrue(resolve_har_files(args))
def test_normal_output_is_concise_summary(self):
result = {
"taskCount": 3,
"balance": {"summary": {"allCash": "56", "totalCoin": 100}},
"actions": [{"signIn": {"ok": True}}, {"watchAd": [{"warnings": ["missing_ad_material"]}]}],
}
with tempfile.TemporaryFile("w+", encoding="utf-8") as fp:
with patch("sys.stdout", fp):
print_normal_output("run-all", result)
fp.seek(0)
text = fp.read()
self.assertIn("任务数量: 3", text)
self.assertIn("执行动作: 2", text)
self.assertIn("余额: 56", text)
self.assertIn("异常摘要", text)
self.assertNotIn("request", text)
if __name__ == "__main__":
unittest.main()