83 lines
3.1 KiB
Python
83 lines
3.1 KiB
Python
import sys
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[2]
|
|
DOMAINCHECK_ROOT = PROJECT_ROOT / "domainCheck"
|
|
if str(DOMAINCHECK_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(DOMAINCHECK_ROOT))
|
|
|
|
import requests
|
|
|
|
from domainCheck.detect import geetest2, jucha, juming, juziseo
|
|
from domainCheck.detect.locked_pickle import load_pickle_locked, save_pickle_atomic
|
|
|
|
|
|
class GeetestCookieSafetyTests(unittest.TestCase):
|
|
def test_geetest_slide_asset_fetch_applies_timeout(self):
|
|
response = MagicMock()
|
|
response.content = b"binary"
|
|
with patch("domainCheck.detect.geetest2.requests.get", return_value=response) as mock_get:
|
|
slider = geetest2.slide()
|
|
with patch.object(slider, "tp_huanyuan", return_value=b"bg-bytes"):
|
|
with patch("domainCheck.detect.geetest2.quekou") as mock_quekou:
|
|
mock_quekou.return_value.get_distance.return_value = 12
|
|
slider.huak({"bg": "bg.png", "slice": "slice.png"})
|
|
|
|
self.assertGreaterEqual(mock_get.call_count, 2)
|
|
for call in mock_get.call_args_list:
|
|
self.assertEqual(slider.asset_timeout, call.kwargs["timeout"])
|
|
|
|
def test_locked_pickle_roundtrip_is_atomic(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
target = Path(temp_dir) / "cookies.pkl"
|
|
save_pickle_atomic(str(target), {"sid": "abc"})
|
|
loaded = load_pickle_locked(str(target), default_factory=dict)
|
|
self.assertEqual({"sid": "abc"}, loaded)
|
|
|
|
def test_juziseo_cookie_roundtrip_uses_locked_pickle(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
target = Path(temp_dir) / "juziseo.pkl"
|
|
detector = juziseo.Juziseo()
|
|
jar = requests.cookies.RequestsCookieJar()
|
|
jar.set("sid", "value")
|
|
detector.cookie = jar
|
|
detector.save_cookies(str(target))
|
|
|
|
loaded = juziseo.Juziseo()
|
|
loaded.load_cookies(str(target))
|
|
|
|
self.assertEqual("value", loaded.cookie.get("sid"))
|
|
|
|
def test_jucha_and_juming_cookie_roundtrip_use_locked_pickle(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
jucha_path = Path(temp_dir) / "jucha.pkl"
|
|
juming_path = Path(temp_dir) / "juming.pkl"
|
|
|
|
jc = jucha.JC()
|
|
jc_jar = requests.cookies.RequestsCookieJar()
|
|
jc_jar.set("jc", "cookie")
|
|
jc.cookie = jc_jar
|
|
jc.save_cookies(str(jucha_path))
|
|
|
|
jm = juming.JM()
|
|
jm_jar = requests.cookies.RequestsCookieJar()
|
|
jm_jar.set("jm", "cookie")
|
|
jm.cookie = jm_jar
|
|
jm.save_cookies(str(juming_path))
|
|
|
|
loaded_jc = jucha.JC()
|
|
loaded_jc.load_cookies(str(jucha_path))
|
|
|
|
loaded_jm = juming.JM()
|
|
loaded_jm.load_cookies(str(juming_path))
|
|
|
|
self.assertEqual("cookie", loaded_jc.cookie.get("jc"))
|
|
self.assertEqual("cookie", loaded_jm.cookie.get("jm"))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|