101 lines
3.7 KiB
Python
101 lines
3.7 KiB
Python
"""Test universal_hunt: _is_junk entropy/repeat/diversity, scan_text, db dedupe.
|
|
|
|
Offline — no network calls.
|
|
"""
|
|
import hashlib
|
|
import os
|
|
import sqlite3
|
|
import sys
|
|
import tempfile
|
|
|
|
# ── setup: temp data dir before importing universal_hunt ──
|
|
_tmp = tempfile.mkdtemp()
|
|
_data = os.path.join(_tmp, "data")
|
|
os.makedirs(_data, exist_ok=True)
|
|
|
|
# Write dummy gh_tokens.txt so PATPool doesn't crash
|
|
with open(os.path.join(_data, "gh_tokens.txt"), "w") as f:
|
|
f.write("ghp_dummy\n")
|
|
|
|
sys.path.insert(0, os.path.abspath(os.path.dirname(__file__)))
|
|
import universal_hunt as uh
|
|
|
|
# Patch globals to use our temp DB
|
|
uh.DB = os.path.join(_data, "_smoke_uni_hunt.db")
|
|
uh.LOG = os.path.join(_data, "_smoke_uni_hunt.log")
|
|
uh.TOKENS_FILE = os.path.join(_data, "gh_tokens.txt")
|
|
|
|
# Recreate conn on our temp DB
|
|
uh.conn = sqlite3.connect(uh.DB, timeout=120, check_same_thread=False,
|
|
isolation_level=None)
|
|
uh.conn.execute("""CREATE TABLE IF NOT EXISTS keys (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
hash TEXT UNIQUE,
|
|
val TEXT, prov TEXT, status TEXT DEFAULT "NEW",
|
|
plan TEXT, price TEXT, remaining TEXT,
|
|
found TEXT, repo TEXT)""")
|
|
|
|
# Reinit pool + reset run state
|
|
uh.POOL = uh.PATPool()
|
|
uh.seen_files.clear()
|
|
uh.seen_raw_hashes.clear()
|
|
uh.raw_fetched[0] = 0
|
|
uh.found_keys.clear()
|
|
|
|
# ── tests ──
|
|
|
|
# 1. _is_junk — low diversity (len(set) < 6)
|
|
assert uh._is_junk("aaaaaabbbbbbccccccdddddd") is True, "len(set)=4 <6"
|
|
assert uh._is_junk("aaaaaaaaaaaaaaaaaaaa") is True, "set=1 <6"
|
|
assert uh._is_junk("aaaabbbbccccddddeee") is True, "set=5 <6"
|
|
|
|
# 2. _is_junk — low entropy for long strings (len>=20, norm_h < 0.2)
|
|
assert uh._is_junk("aaaaaaaaaaaaaaaaaaaabbbbbbbbbbbbbbbbbbbb") is True, \
|
|
"low-entropy 40-char"
|
|
# High-entropy string should pass
|
|
assert uh._is_junk("K8x9mP2qR5vWnJ3cL7f") is False, "high-entropy short"
|
|
assert uh._is_junk("abc") is True, "short set<6"
|
|
|
|
# 3. _is_junk — repeating pattern (blk 1-4)
|
|
assert uh._is_junk("abababababababababab") is True, "repeating ab"
|
|
assert uh._is_junk("xyzxyzxyzxyzxyzxyzxy") is True, "repeating xyz (plen=3)"
|
|
assert uh._is_junk("abcabcabcabcabcabcab") is True, "repeating abc"
|
|
|
|
# 4. Non-repeating, high entropy, set>=6 — should pass
|
|
assert uh._is_junk("sk-proj-B1cD2eF3gH4iJ5kL6mN7oP8qR9sT0uV1wX2yZ3aB4cD5eF6gH7iJ8kL9mN0oP") is False, \
|
|
"real-looking key"
|
|
|
|
# 5. scan_text — GARBAGE filtering (separate from _is_junk)
|
|
junk_fixture = "OPENAI_API_KEY=sk-proj-xxxx-test-key-dummy-placeholder"
|
|
junk_hits = list(uh.scan_text(junk_fixture))
|
|
assert len(junk_hits) == 0, f"expected 0 junk hits got {len(junk_hits)}"
|
|
|
|
# 8. db_upsert — writes and dedupes
|
|
h1 = uh.db_upsert("sk-proj-TESTKEY-ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789abcdefghij",
|
|
"OPENAI", "test/repo")
|
|
h2 = uh.db_upsert("sk-proj-TESTKEY-OTHERABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789abcdef",
|
|
"OPENAI", "test/repo2")
|
|
assert h1 != h2, "different keys -> different hash"
|
|
|
|
# Dedupe: same key writes once
|
|
h1a = uh.db_upsert("sk-proj-TESTKEY-ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789abcdefghij",
|
|
"OPENAI", "test/repo")
|
|
assert h1a == h1, "same key -> same hash"
|
|
|
|
row = uh.conn.execute("SELECT count(*) FROM keys").fetchone()[0]
|
|
assert row == 2, f"expected 2 rows got {row}"
|
|
|
|
# 9. scan_text — openai_proj normalizes to OPENAI
|
|
norm_text = "VITE_OPENAI_API_KEY=sk-proj-B1cD2eF3gH4iJ5kL6mN7oP8qR9sT0uV1wX2yZ3aB4cD5eF6gH7iJ8kL9mN0oP"
|
|
norm_hits = list(uh.scan_text(norm_text))
|
|
has_openai = any(p == "OPENAI" for p, _ in norm_hits)
|
|
assert has_openai, f"normalization failed: {norm_hits}"
|
|
|
|
# 10. trust_env is False
|
|
assert uh.S.trust_env is False, "trust_env must be False"
|
|
|
|
# 11. seen_raw_hashes exists and works
|
|
uh.seen_raw_hashes.add("deadbeef")
|
|
assert "deadbeef" in uh.seen_raw_hashes, "seen_raw_hashes add"
|
|
|
|
print("ALL PASS") |