236 lines
10 KiB
Python
236 lines
10 KiB
Python
"""API keys 多 key 单测：primary + 附加 key、撤销、revoke-others、
|
||
登录签发附加 key（不再轮换）。
|
||
|
||
测试里出现的 key 一律是测试内生成的假 key，绝无真实 secret。
|
||
"""
|
||
|
||
import hashlib
|
||
import http.client
|
||
import json
|
||
import os
|
||
import sys
|
||
import tempfile
|
||
import threading
|
||
import unittest
|
||
from pathlib import Path
|
||
|
||
try:
|
||
from cloud.auth import AuthStore
|
||
from cloud.oauth import map_external_identity_to_local_user_with_key
|
||
from cloud.server import create_server
|
||
except ImportError:
|
||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||
from cloud.auth import AuthStore
|
||
from cloud.oauth import map_external_identity_to_local_user_with_key
|
||
from cloud.server import create_server
|
||
|
||
|
||
class ApiKeysFunctionTestCase(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.db_path = os.path.join(self.tmp.name, "control.db")
|
||
self.auth = AuthStore(self.db_path)
|
||
self.user_id, self.primary = self.auth.create_user("user-a")
|
||
|
||
def tearDown(self):
|
||
try:
|
||
self.auth.conn.close()
|
||
except Exception:
|
||
pass
|
||
self.tmp.cleanup()
|
||
|
||
def test_primary_listed_as_legacy(self):
|
||
keys = self.auth.list_api_keys(self.user_id)
|
||
self.assertEqual(len(keys), 1)
|
||
self.assertEqual(keys[0]["id"], "legacy")
|
||
self.assertEqual(keys[0]["label"], "primary")
|
||
self.assertIsNone(keys[0]["hint"])
|
||
self.assertIsNone(keys[0]["revoked_at"])
|
||
# 列表里绝无明文与哈希。
|
||
for key in keys:
|
||
self.assertNotIn("api_key", key)
|
||
self.assertNotIn("key_hash", key)
|
||
|
||
def test_additional_key_authenticates_and_revokes(self):
|
||
key_id, plain = self.auth.create_api_key(self.user_id, "laptop")
|
||
self.assertEqual(self.auth.authenticate(plain), self.user_id)
|
||
self.assertEqual(plain[-4:], [
|
||
k for k in self.auth.list_api_keys(self.user_id)
|
||
if k["id"] == key_id][0]["hint"])
|
||
self.assertTrue(self.auth.revoke_api_key(self.user_id, key_id))
|
||
self.assertIsNone(self.auth.authenticate(plain))
|
||
self.assertEqual(self.auth.authenticate(self.primary), self.user_id)
|
||
self.assertFalse(self.auth.revoke_api_key(self.user_id, key_id))
|
||
|
||
def test_revoke_legacy_keeps_additional(self):
|
||
_key_id, plain = self.auth.create_api_key(self.user_id, "ci")
|
||
self.assertTrue(self.auth.revoke_api_key(self.user_id, "legacy"))
|
||
self.assertIsNone(self.auth.authenticate(self.primary))
|
||
self.assertEqual(self.auth.authenticate(plain), self.user_id)
|
||
self.assertFalse(self.auth.revoke_api_key(self.user_id, "legacy"))
|
||
|
||
def test_revoke_others_keeps_only_current(self):
|
||
_id1, plain1 = self.auth.create_api_key(self.user_id, "a")
|
||
_id2, plain2 = self.auth.create_api_key(self.user_id, "b")
|
||
keep = hashlib.sha256(plain2.encode()).hexdigest()
|
||
count = self.auth.revoke_other_api_keys(self.user_id, keep)
|
||
self.assertEqual(count, 2) # primary + 第一个附加 key
|
||
self.assertIsNone(self.auth.authenticate(self.primary))
|
||
self.assertIsNone(self.auth.authenticate(plain1))
|
||
self.assertEqual(self.auth.authenticate(plain2), self.user_id)
|
||
|
||
def test_last_used_throttled(self):
|
||
key_id, plain = self.auth.create_api_key(self.user_id, "x")
|
||
self.assertEqual(self.auth.authenticate(plain), self.user_id)
|
||
first = [k for k in self.auth.list_api_keys(self.user_id)
|
||
if k["id"] == key_id][0]["last_used_at"]
|
||
self.assertIsNotNone(first)
|
||
self.assertEqual(self.auth.authenticate(plain), self.user_id)
|
||
second = [k for k in self.auth.list_api_keys(self.user_id)
|
||
if k["id"] == key_id][0]["last_used_at"]
|
||
self.assertEqual(first, second) # 一小时内不重复写库
|
||
|
||
def test_create_for_unknown_user(self):
|
||
with self.assertRaises(ValueError):
|
||
self.auth.create_api_key("no-such-user", "x")
|
||
|
||
def test_oauth_relogin_issues_additional_not_rotation(self):
|
||
# 已存在 oauth 用户再次登录：签发附加 key，旧 primary 不死。
|
||
user_id, primary = self.auth.create_user("oauth:github:gh-user")
|
||
_uid, created, plain, rotated = \
|
||
map_external_identity_to_local_user_with_key(
|
||
self.auth, "github", {"login": "gh-user",
|
||
"external_id": "4242"})
|
||
self.assertEqual(_uid, user_id)
|
||
self.assertFalse(created)
|
||
self.assertFalse(rotated)
|
||
self.assertNotEqual(plain, primary)
|
||
self.assertEqual(self.auth.authenticate(primary), user_id)
|
||
self.assertEqual(self.auth.authenticate(plain), user_id)
|
||
|
||
|
||
class ApiKeysHttpTestCase(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.db_path = os.path.join(self.tmp.name, "control.db")
|
||
self.auth = AuthStore(self.db_path)
|
||
self.user_a, self.key_a = self.auth.create_user("user-a")
|
||
self.server = create_server(self.db_path, "127.0.0.1", 0)
|
||
self.port = self.server.server_address[1]
|
||
self.thread = threading.Thread(
|
||
target=self.server.serve_forever, daemon=True)
|
||
self.thread.start()
|
||
|
||
def tearDown(self):
|
||
self.server.shutdown()
|
||
self.server.server_close()
|
||
self.thread.join(timeout=5)
|
||
try:
|
||
self.auth.conn.close()
|
||
except Exception:
|
||
pass
|
||
self.tmp.cleanup()
|
||
|
||
def _req(self, method, path, body=None, key=None):
|
||
conn = http.client.HTTPConnection("127.0.0.1", self.port, timeout=10)
|
||
try:
|
||
headers = {}
|
||
payload = None
|
||
if body is not None:
|
||
payload = json.dumps(body)
|
||
headers["Content-Type"] = "application/json"
|
||
if key is not None:
|
||
headers["X-API-Key"] = key
|
||
conn.request(method, path, body=payload, headers=headers)
|
||
resp = conn.getresponse()
|
||
raw = resp.read().decode()
|
||
try:
|
||
parsed = json.loads(raw)
|
||
except ValueError:
|
||
parsed = raw
|
||
return resp.status, parsed
|
||
finally:
|
||
conn.close()
|
||
|
||
def test_keys_require_auth(self):
|
||
status, _ = self._req("GET", "/users/me/keys")
|
||
self.assertEqual(status, 401)
|
||
status, _ = self._req("POST", "/users/me/keys", {})
|
||
self.assertEqual(status, 401)
|
||
|
||
def test_create_list_revoke_flow(self):
|
||
status, body = self._req("GET", "/users/me/keys", key=self.key_a)
|
||
self.assertEqual(status, 200, body)
|
||
self.assertEqual(len(body["keys"]), 1)
|
||
self.assertEqual(body["keys"][0]["id"], "legacy")
|
||
# 创建附加 key：明文只出现一次，hint 为后 4 位。
|
||
status, body = self._req(
|
||
"POST", "/users/me/keys", {"label": "laptop"}, key=self.key_a)
|
||
self.assertEqual(status, 201, body)
|
||
new_key = body["api_key"]
|
||
self.assertEqual(body["hint"], new_key[-4:])
|
||
key_id = body["id"]
|
||
# 新 key 可用；列表里有两条且无明文。
|
||
status, _ = self._req("GET", "/users/me/keys", key=new_key)
|
||
self.assertEqual(status, 200)
|
||
status, body = self._req("GET", "/users/me/keys", key=self.key_a)
|
||
self.assertEqual(len(body["keys"]), 2)
|
||
self.assertNotIn(new_key, json.dumps(body))
|
||
# 撤销附加 key：新 key 401，primary 仍可用；再撤 404。
|
||
status, body = self._req(
|
||
"DELETE", "/users/me/keys/" + key_id, key=self.key_a)
|
||
self.assertEqual(status, 200, body)
|
||
status, _ = self._req("GET", "/users/me/keys", key=new_key)
|
||
self.assertEqual(status, 401)
|
||
status, _ = self._req("GET", "/users/me/keys", key=self.key_a)
|
||
self.assertEqual(status, 200)
|
||
status, _ = self._req(
|
||
"DELETE", "/users/me/keys/" + key_id, key=self.key_a)
|
||
self.assertEqual(status, 404)
|
||
# 撤销 primary（legacy）：primary 401；再撤需用另一把 key，
|
||
# 结果应为 404（已撤销），且另一把 key 不受影响。
|
||
status, body = self._req(
|
||
"POST", "/users/me/keys", {"label": "spare"}, key=self.key_a)
|
||
self.assertEqual(status, 201, body)
|
||
spare_key = body["api_key"]
|
||
status, body = self._req(
|
||
"DELETE", "/users/me/keys/legacy", key=self.key_a)
|
||
self.assertEqual(status, 200, body)
|
||
status, _ = self._req("GET", "/users/me/keys", key=self.key_a)
|
||
self.assertEqual(status, 401)
|
||
status, _ = self._req(
|
||
"DELETE", "/users/me/keys/legacy", key=spare_key)
|
||
self.assertEqual(status, 404)
|
||
status, _ = self._req("GET", "/users/me/keys", key=spare_key)
|
||
self.assertEqual(status, 200)
|
||
|
||
def test_revoke_legacy_and_revoke_others(self):
|
||
_s, created = self._req(
|
||
"POST", "/users/me/keys", {"label": "ci"}, key=self.key_a)
|
||
ci_key = created["api_key"]
|
||
_s, created2 = self._req(
|
||
"POST", "/users/me/keys", {"label": "home"}, key=self.key_a)
|
||
home_key = created2["api_key"]
|
||
# revoke-others（用 home key 调）：primary + ci 被撤销。
|
||
status, body = self._req(
|
||
"POST", "/users/me/keys/revoke-others", {}, key=home_key)
|
||
self.assertEqual(status, 200, body)
|
||
self.assertEqual(body["revoked_count"], 2)
|
||
status, _ = self._req("GET", "/users/me/keys", key=self.key_a)
|
||
self.assertEqual(status, 401)
|
||
status, _ = self._req("GET", "/users/me/keys", key=ci_key)
|
||
self.assertEqual(status, 401)
|
||
status, _ = self._req("GET", "/users/me/keys", key=home_key)
|
||
self.assertEqual(status, 200)
|
||
# primary 已被 revoke-others 撤销，再撤 legacy 应为 404；
|
||
# home key 本身不受影响。
|
||
status, body = self._req(
|
||
"DELETE", "/users/me/keys/legacy", key=home_key)
|
||
self.assertEqual(status, 404, body)
|
||
status, _ = self._req("GET", "/users/me/keys", key=home_key)
|
||
self.assertEqual(status, 200)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|