350 lines
14 KiB
Python
350 lines
14 KiB
Python
"""Encrypted on-disk store for OAuth clients, sessions and token hashes.
|
|
|
|
One file, ``DATA_DIR/oauth-store.json``, holds the whole state as a single
|
|
Fernet-encrypted JSON document. The Fernet key is derived from
|
|
``TOKEN_ENCRYPTION_KEY`` with HKDF-SHA256, so the setting can be any long random
|
|
string (the setup command writes a ``secrets.token_urlsafe`` value, not a Fernet
|
|
key).
|
|
|
|
What is stored:
|
|
|
|
* registered MCP clients, including their client secrets (the SDK compares the
|
|
presented secret with the stored one at ``/token``);
|
|
* sessions: the Gitea login and the user's Gitea access and refresh tokens;
|
|
* our own access and refresh tokens, only as SHA-256 hashes.
|
|
|
|
File safety: the directory is created with mode 0700 (an existing one must be owned by the
|
|
running user and is tightened to 0700, otherwise startup stops) and the file is written
|
|
with mode 0600 through a fresh ``O_EXCL`` temporary file and an atomic
|
|
``os.replace``. A symlink at the store path is refused. A file that cannot be
|
|
decrypted or parsed stops startup with a clear message instead of being wiped,
|
|
because wiping would silently sign every user out and hide a wrong key.
|
|
|
|
Mutations happen synchronously in memory (so a check-and-delete is atomic under
|
|
asyncio) and are persisted with :meth:`OAuthStore.save`.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import base64
|
|
import contextlib
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
import os
|
|
import secrets
|
|
import stat
|
|
import time
|
|
from collections.abc import Callable, Iterator
|
|
from dataclasses import asdict, dataclass, field
|
|
from pathlib import Path
|
|
from typing import Any, Literal
|
|
|
|
import anyio.to_thread
|
|
from cryptography.fernet import Fernet, InvalidToken
|
|
from cryptography.hazmat.primitives import hashes
|
|
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
STORE_FILENAME = "oauth-store.json"
|
|
STORE_VERSION = 1
|
|
HKDF_INFO = b"aegis-gitea-mcp oauth store"
|
|
|
|
#: A session that never received tokens (an abandoned sign-in) is dropped after this long.
|
|
ABANDONED_SESSION_SECONDS = 600.0
|
|
|
|
Clock = Callable[[], float]
|
|
TokenKind = Literal["access", "refresh"]
|
|
|
|
|
|
class OAuthStoreError(Exception):
|
|
"""The OAuth store cannot be used (startup error; the message is safe to print)."""
|
|
|
|
|
|
def derive_fernet_key(secret: str) -> bytes:
|
|
"""Derive the Fernet key for the store from ``TOKEN_ENCRYPTION_KEY`` (HKDF-SHA256)."""
|
|
hkdf = HKDF(algorithm=hashes.SHA256(), length=32, salt=None, info=HKDF_INFO)
|
|
return base64.urlsafe_b64encode(hkdf.derive(secret.encode("utf-8")))
|
|
|
|
|
|
def token_hash(value: str) -> str:
|
|
"""SHA-256 hex digest used as the stored form of our tokens, codes and states."""
|
|
return hashlib.sha256(value.encode("utf-8")).hexdigest()
|
|
|
|
|
|
@dataclass
|
|
class ClientRecord:
|
|
"""A dynamically registered MCP client.
|
|
|
|
Attributes:
|
|
info: ``OAuthClientInformationFull`` as JSON (includes the client secret).
|
|
registered_at: Registration time (epoch seconds).
|
|
last_used_at: Last time the client was looked up (epoch seconds).
|
|
"""
|
|
|
|
info: dict[str, Any] = field(repr=False)
|
|
registered_at: float
|
|
last_used_at: float
|
|
|
|
|
|
@dataclass
|
|
class SessionRecord:
|
|
"""One signed-in grant: a Gitea user, the client it was granted to and the Gitea tokens.
|
|
|
|
Attributes:
|
|
session_id: Random identifier (never sent to clients).
|
|
client_id: MCP client the grant belongs to.
|
|
login: Gitea login of the user.
|
|
is_admin: Whether Gitea reported the user as site admin at sign-in
|
|
(informational only; Gitea enforces permissions).
|
|
gitea_access_token: The user's Gitea access token.
|
|
gitea_refresh_token: The user's Gitea refresh token, if Gitea issued one.
|
|
gitea_expires_at: Expiry of the Gitea access token (epoch seconds), if known.
|
|
scopes: Our scopes granted to the client.
|
|
resource: The resource (MCP endpoint URL) the grant is bound to.
|
|
created_at: Sign-in time (epoch seconds).
|
|
previous_refresh_hash: Hash of the last rotated-out refresh token (reuse detection).
|
|
"""
|
|
|
|
session_id: str
|
|
client_id: str
|
|
login: str
|
|
is_admin: bool
|
|
gitea_access_token: str = field(repr=False)
|
|
gitea_refresh_token: str | None = field(repr=False)
|
|
gitea_expires_at: float | None
|
|
scopes: list[str]
|
|
resource: str
|
|
created_at: float
|
|
previous_refresh_hash: str | None = field(default=None, repr=False)
|
|
|
|
|
|
@dataclass
|
|
class TokenRecord:
|
|
"""One of our tokens, stored under its SHA-256 hash."""
|
|
|
|
kind: TokenKind
|
|
session_id: str
|
|
client_id: str
|
|
scopes: list[str]
|
|
expires_at: float
|
|
|
|
|
|
def _restrict_directory(directory: Path) -> None:
|
|
"""Create the state directory with mode 0700, or check and tighten an existing one.
|
|
|
|
Whoever controls the directory can swap the store for an older copy (which
|
|
resurrects revoked grants) or delete it, so an existing directory must belong
|
|
to the effective user and must not stay writable by group or others.
|
|
|
|
Raises:
|
|
OAuthStoreError: The directory is owned by another user or remains writable
|
|
by group or others after an attempt to tighten it.
|
|
"""
|
|
if not directory.exists():
|
|
directory.mkdir(mode=0o700, parents=True, exist_ok=True)
|
|
return
|
|
info = os.stat(directory)
|
|
if not stat.S_ISDIR(info.st_mode):
|
|
raise OAuthStoreError(f"{directory} is not a directory")
|
|
owner = getattr(os, "geteuid", lambda: info.st_uid)()
|
|
if info.st_uid != owner:
|
|
raise OAuthStoreError(
|
|
f"{directory} is owned by another user (uid {info.st_uid}, this process runs as "
|
|
f"uid {owner}); refusing to keep the OAuth token store there. Use a DATA_DIR "
|
|
"owned by the service user"
|
|
)
|
|
mode = stat.S_IMODE(info.st_mode)
|
|
if mode & 0o077:
|
|
try:
|
|
directory.chmod(0o700)
|
|
except OSError:
|
|
logger.warning("could not restrict the permissions of the state directory to 0700")
|
|
if stat.S_IMODE(os.stat(directory).st_mode) & 0o022:
|
|
raise OAuthStoreError(
|
|
f"{directory} is writable by group or others and could not be restricted; "
|
|
"refusing to keep the OAuth token store there. Run 'chmod 700' on it"
|
|
)
|
|
|
|
|
|
def _parse_document(document: Any) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]:
|
|
if not isinstance(document, dict) or document.get("version") != STORE_VERSION:
|
|
raise ValueError("unsupported store version")
|
|
clients_raw = document.get("clients", {})
|
|
sessions_raw = document.get("sessions", {})
|
|
tokens_raw = document.get("tokens", {})
|
|
if not all(isinstance(part, dict) for part in (clients_raw, sessions_raw, tokens_raw)):
|
|
raise ValueError("malformed store sections")
|
|
clients = {str(key): ClientRecord(**value) for key, value in clients_raw.items()}
|
|
sessions = {str(key): SessionRecord(**value) for key, value in sessions_raw.items()}
|
|
tokens = {str(key): TokenRecord(**value) for key, value in tokens_raw.items()}
|
|
for record in clients.values():
|
|
if not isinstance(record.info, dict):
|
|
raise ValueError("malformed client record")
|
|
for token in tokens.values():
|
|
if token.kind not in ("access", "refresh"):
|
|
raise ValueError("malformed token record")
|
|
return clients, sessions, tokens
|
|
|
|
|
|
class OAuthStore:
|
|
"""In-memory OAuth state backed by one encrypted file.
|
|
|
|
Use :meth:`load` to create it. All getters and mutators are synchronous;
|
|
:meth:`save` persists the current state (pruning expired entries first).
|
|
|
|
Args:
|
|
directory: State directory (``DATA_DIR``).
|
|
secret: ``TOKEN_ENCRYPTION_KEY``.
|
|
clock: Wall clock (epoch seconds); tests inject a fake one.
|
|
"""
|
|
|
|
def __init__(self, directory: Path, secret: str, clock: Clock = time.time) -> None:
|
|
self._directory = directory
|
|
self._path = directory / STORE_FILENAME
|
|
self._fernet = Fernet(derive_fernet_key(secret))
|
|
self._clock = clock
|
|
self._lock = asyncio.Lock()
|
|
self.clients: dict[str, ClientRecord] = {}
|
|
self.sessions: dict[str, SessionRecord] = {}
|
|
self.tokens: dict[str, TokenRecord] = {}
|
|
|
|
@classmethod
|
|
def load(cls, directory: Path, secret: str, clock: Clock = time.time) -> OAuthStore:
|
|
"""Open the store, reading and decrypting the file when it exists.
|
|
|
|
Raises:
|
|
OAuthStoreError: The directory belongs to another user or stays group or
|
|
world writable, the path is a symlink or not a regular file, or the
|
|
file cannot be decrypted or parsed.
|
|
"""
|
|
store = cls(directory, secret, clock)
|
|
_restrict_directory(directory)
|
|
store._read()
|
|
store.prune()
|
|
return store
|
|
|
|
@property
|
|
def path(self) -> Path:
|
|
"""Location of the store file."""
|
|
return self._path
|
|
|
|
# Reading and writing the file
|
|
|
|
def _read(self) -> None:
|
|
try:
|
|
info = os.lstat(self._path)
|
|
except FileNotFoundError:
|
|
return
|
|
if stat.S_ISLNK(info.st_mode):
|
|
raise OAuthStoreError(
|
|
f"{self._path} is a symlink; refusing to use it as the OAuth token store"
|
|
)
|
|
if not stat.S_ISREG(info.st_mode):
|
|
raise OAuthStoreError(f"{self._path} is not a regular file")
|
|
fd = os.open(self._path, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0))
|
|
with os.fdopen(fd, "rb") as handle:
|
|
raw = handle.read()
|
|
if stat.S_IMODE(info.st_mode) & 0o077:
|
|
logger.warning("OAuth store file had loose permissions; restricting it to 0600")
|
|
os.chmod(self._path, 0o600)
|
|
try:
|
|
plain = self._fernet.decrypt(raw)
|
|
except InvalidToken:
|
|
raise OAuthStoreError(
|
|
f"cannot decrypt {self._path}: TOKEN_ENCRYPTION_KEY does not match the key the "
|
|
"store was written with, or the file is damaged. Restore the original key, or "
|
|
"move the file away to start over (this signs every OAuth user out)."
|
|
) from None
|
|
try:
|
|
self.clients, self.sessions, self.tokens = _parse_document(json.loads(plain))
|
|
except (ValueError, TypeError, KeyError):
|
|
raise OAuthStoreError(
|
|
f"{self._path} is not a valid OAuth store document; move the file away to "
|
|
"start over (this signs every OAuth user out)"
|
|
) from None
|
|
|
|
def _document(self) -> bytes:
|
|
document = {
|
|
"version": STORE_VERSION,
|
|
"clients": {key: asdict(value) for key, value in self.clients.items()},
|
|
"sessions": {key: asdict(value) for key, value in self.sessions.items()},
|
|
"tokens": {key: asdict(value) for key, value in self.tokens.items()},
|
|
}
|
|
return json.dumps(document, separators=(",", ":")).encode("utf-8")
|
|
|
|
def _write(self, data: bytes) -> None:
|
|
_restrict_directory(self._directory)
|
|
with contextlib.suppress(FileNotFoundError):
|
|
if stat.S_ISLNK(os.lstat(self._path).st_mode):
|
|
raise OAuthStoreError(
|
|
f"{self._path} is a symlink; refusing to write the OAuth token store"
|
|
)
|
|
temporary = self._directory / f".{STORE_FILENAME}.{secrets.token_hex(8)}.tmp"
|
|
flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0)
|
|
fd = os.open(temporary, flags, 0o600)
|
|
try:
|
|
with os.fdopen(fd, "wb") as handle:
|
|
handle.write(data)
|
|
handle.flush()
|
|
os.fsync(handle.fileno())
|
|
os.replace(temporary, self._path)
|
|
except BaseException:
|
|
with contextlib.suppress(FileNotFoundError):
|
|
os.unlink(temporary)
|
|
raise
|
|
|
|
async def save(self) -> None:
|
|
"""Prune expired entries and persist the state atomically (encrypted, mode 0600)."""
|
|
async with self._lock:
|
|
self.prune()
|
|
data = self._fernet.encrypt(self._document())
|
|
# The write fsyncs: keep that off the event loop (the bytes are already built,
|
|
# so the thread touches no shared state).
|
|
await anyio.to_thread.run_sync(self._write, data)
|
|
|
|
# Maintenance
|
|
|
|
def prune(self) -> None:
|
|
"""Drop expired tokens, and sessions that have no tokens left."""
|
|
now = self._clock()
|
|
for key in [key for key, token in self.tokens.items() if token.expires_at <= now]:
|
|
del self.tokens[key]
|
|
live = {token.session_id for token in self.tokens.values()}
|
|
for session_id, session in list(self.sessions.items()):
|
|
if session_id in live:
|
|
continue
|
|
if session.created_at + ABANDONED_SESSION_SECONDS < now:
|
|
del self.sessions[session_id]
|
|
|
|
# Sessions and tokens
|
|
|
|
def session_token_hashes(self, session_id: str) -> list[str]:
|
|
"""Hashes of every token issued for ``session_id``."""
|
|
return [key for key, token in self.tokens.items() if token.session_id == session_id]
|
|
|
|
def end_session(self, session_id: str) -> bool:
|
|
"""Delete a session and all of its tokens; return whether it existed."""
|
|
for key in self.session_token_hashes(session_id):
|
|
del self.tokens[key]
|
|
return self.sessions.pop(session_id, None) is not None
|
|
|
|
def sessions_of(self, login: str) -> Iterator[SessionRecord]:
|
|
"""Every session of one Gitea login."""
|
|
return (session for session in self.sessions.values() if session.login == login)
|
|
|
|
def session_by_previous_refresh(self, refresh_hash: str) -> SessionRecord | None:
|
|
"""The session whose last rotated-out refresh token has ``refresh_hash``."""
|
|
for session in self.sessions.values():
|
|
if session.previous_refresh_hash is not None and secrets.compare_digest(
|
|
session.previous_refresh_hash, refresh_hash
|
|
):
|
|
return session
|
|
return None
|
|
|
|
def client_has_sessions(self, client_id: str) -> bool:
|
|
"""Whether any session belongs to ``client_id``."""
|
|
return any(session.client_id == client_id for session in self.sessions.values())
|