68 lines
2.1 KiB
Python
68 lines
2.1 KiB
Python
"""Controlled local PostgreSQL connector for the standalone MCP process."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import stat
|
|
from pathlib import Path
|
|
from typing import Any, Callable, Dict
|
|
|
|
|
|
_REQUIRED = (
|
|
"ARR_DB_HOST",
|
|
"ARR_DB_PORT",
|
|
"ARR_DB_USER",
|
|
"ARR_DB_PASSWORD",
|
|
"ARR_DB_NAME",
|
|
)
|
|
|
|
|
|
def _read_private_database_config(path: Path) -> Dict[str, object]:
|
|
candidate = path.expanduser()
|
|
if candidate.is_symlink():
|
|
raise ValueError("controlled database configuration is invalid")
|
|
resolved = candidate.resolve(strict=True)
|
|
metadata = resolved.stat()
|
|
if (
|
|
not stat.S_ISREG(metadata.st_mode)
|
|
or metadata.st_uid != os.getuid()
|
|
or stat.S_IMODE(metadata.st_mode) != 0o600
|
|
):
|
|
raise ValueError("controlled database configuration is invalid")
|
|
values: Dict[str, str] = {}
|
|
for raw in resolved.read_text(encoding="utf-8").splitlines():
|
|
line = raw.strip()
|
|
if not line or line.startswith("#") or "=" not in line:
|
|
continue
|
|
key, value = line.split("=", 1)
|
|
key = key.strip()
|
|
if key in values:
|
|
raise ValueError("controlled database configuration is invalid")
|
|
values[key] = value.strip().strip('"').strip("'")
|
|
if any(not values.get(key) for key in _REQUIRED):
|
|
raise ValueError("controlled database configuration is incomplete")
|
|
try:
|
|
port = int(values["ARR_DB_PORT"])
|
|
except ValueError:
|
|
raise ValueError("controlled database configuration is invalid") from None
|
|
if not 1 <= port <= 65535:
|
|
raise ValueError("controlled database configuration is invalid")
|
|
return {
|
|
"host": values["ARR_DB_HOST"],
|
|
"port": port,
|
|
"user": values["ARR_DB_USER"],
|
|
"password": values["ARR_DB_PASSWORD"],
|
|
"dbname": values["ARR_DB_NAME"],
|
|
}
|
|
|
|
|
|
def controlled_connect(path: Path) -> Callable[[str], Any]:
|
|
import psycopg # type: ignore[import-not-found]
|
|
|
|
parameters = _read_private_database_config(path)
|
|
|
|
def connect(_dsn: str) -> Any:
|
|
return psycopg.connect(**parameters, autocommit=False)
|
|
|
|
return connect
|