150 lines
5.0 KiB
Python
150 lines
5.0 KiB
Python
from __future__ import annotations
|
|
|
|
import io
|
|
import json
|
|
import os
|
|
import tempfile
|
|
import unittest
|
|
import urllib.error
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
from scripts.revoke_demo_token import revoke
|
|
from scripts.verify_demo_token_flow import post_tool_call, verify_flow
|
|
|
|
|
|
class RevokeDemoTokenScriptTests(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self.temp_dir = tempfile.TemporaryDirectory()
|
|
self.path = Path(self.temp_dir.name) / "revocations.json"
|
|
|
|
def tearDown(self) -> None:
|
|
self.temp_dir.cleanup()
|
|
|
|
def test_creates_file_when_missing(self) -> None:
|
|
jti = "a" * 32
|
|
result = revoke(self.path, jti)
|
|
self.assertEqual(result, [jti])
|
|
document = json.loads(self.path.read_text(encoding="utf-8"))
|
|
self.assertEqual(document["revokedJtis"], [jti])
|
|
if os.name != "nt":
|
|
self.assertEqual(os.stat(self.path).st_mode & 0o777, 0o600)
|
|
|
|
def test_appends_to_existing_file(self) -> None:
|
|
self.path.write_text(
|
|
json.dumps({"schemaVersion": "1.0", "revokedJtis": ["a" * 32]}),
|
|
encoding="utf-8",
|
|
)
|
|
original = os.stat(self.path)
|
|
result = revoke(self.path, "b" * 32)
|
|
self.assertEqual(sorted(result), ["a" * 32, "b" * 32])
|
|
if os.name != "nt":
|
|
updated = os.stat(self.path)
|
|
self.assertEqual((updated.st_uid, updated.st_gid), (original.st_uid, original.st_gid))
|
|
|
|
def test_revoking_same_jti_twice_is_idempotent(self) -> None:
|
|
revoke(self.path, "a" * 32)
|
|
result = revoke(self.path, "a" * 32)
|
|
self.assertEqual(result, ["a" * 32])
|
|
|
|
def test_rejects_invalid_jti_and_leaves_existing_file_unchanged(self) -> None:
|
|
self.path.write_text(
|
|
json.dumps({"schemaVersion": "1.0", "revokedJtis": ["a" * 32]}),
|
|
encoding="utf-8",
|
|
)
|
|
original = self.path.read_bytes()
|
|
with self.assertRaises(ValueError):
|
|
revoke(self.path, "not-a-jti")
|
|
self.assertEqual(self.path.read_bytes(), original)
|
|
|
|
|
|
class DemoTokenFlowScriptTests(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self.temp_dir = tempfile.TemporaryDirectory()
|
|
self.root = Path(self.temp_dir.name)
|
|
self.key_file = self.root / "demo_token_signing_key"
|
|
self.key_file.write_text("x" * 32, encoding="utf-8")
|
|
if os.name != "nt":
|
|
self.key_file.chmod(0o600)
|
|
self.revocations = self.root / "revocations.json"
|
|
|
|
def tearDown(self) -> None:
|
|
self.temp_dir.cleanup()
|
|
|
|
def test_validates_json_then_revokes_the_generated_token(self) -> None:
|
|
class Headers:
|
|
def get_content_type(self) -> str:
|
|
return "application/json"
|
|
|
|
class SuccessResponse:
|
|
status = 200
|
|
headers = Headers()
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *_args):
|
|
return False
|
|
|
|
def read(self) -> bytes:
|
|
return b'{"jsonrpc":"2.0","result":{}}'
|
|
|
|
calls = 0
|
|
|
|
def fake_urlopen(request, *, timeout):
|
|
nonlocal calls
|
|
calls += 1
|
|
self.assertGreater(timeout, 0)
|
|
self.assertEqual(request.get_header("Authorization", "")[:12], "Bearer demo_")
|
|
if calls == 1:
|
|
return SuccessResponse()
|
|
raise urllib.error.HTTPError(
|
|
request.full_url,
|
|
401,
|
|
"Unauthorized",
|
|
Headers(),
|
|
io.BytesIO(b"{}"),
|
|
)
|
|
|
|
with patch("urllib.request.urlopen", side_effect=fake_urlopen):
|
|
valid, revoked = verify_flow(
|
|
"http://127.0.0.1/mcp/mobilityops/prod",
|
|
self.key_file,
|
|
self.revocations,
|
|
timeout_seconds=2,
|
|
)
|
|
|
|
self.assertEqual(valid.status, 200)
|
|
self.assertEqual(revoked.status, 401)
|
|
self.assertEqual(calls, 2)
|
|
document = json.loads(self.revocations.read_text(encoding="utf-8"))
|
|
self.assertEqual(len(document["revokedJtis"]), 1)
|
|
|
|
def test_http_error_with_reset_body_still_returns_bounded_status(self) -> None:
|
|
class ResetHeaders:
|
|
def get_content_type(self) -> str:
|
|
return "application/json"
|
|
|
|
error = urllib.error.HTTPError(
|
|
"http://127.0.0.1/mcp",
|
|
401,
|
|
"Unauthorized",
|
|
ResetHeaders(),
|
|
None,
|
|
)
|
|
|
|
def broken_read(*_args, **_kwargs):
|
|
raise ConnectionResetError(104, "Connection reset by peer")
|
|
|
|
error.read = broken_read # type: ignore[method-assign]
|
|
with patch("urllib.request.urlopen", side_effect=error):
|
|
result = post_tool_call(
|
|
"http://127.0.0.1/mcp",
|
|
"demo_test",
|
|
timeout_seconds=2,
|
|
)
|
|
|
|
self.assertEqual(result.status, 401)
|
|
self.assertEqual(result.content_type, "application/json")
|
|
self.assertEqual(result.body, b"")
|