146 lines
7.7 KiB
Python
146 lines
7.7 KiB
Python
"""DB-Dienst für PostgreSQL: Tools, nur lesen, Zugangsdaten, Anmeldung, Audit-Log."""
|
||
|
||
import json
|
||
import logging
|
||
|
||
import httpx2
|
||
import pytest
|
||
|
||
from dbmcp.core import Invalid
|
||
from dbmcp.core.config import Settings
|
||
from dbmcp.core.connection import from_headers
|
||
from dbmcp.core.sql import check_query
|
||
from tests.conftest import KEY, PASSWORD
|
||
|
||
|
||
def text(result) -> str:
|
||
return "".join(c.text for c in result.content if c.type == "text")
|
||
|
||
|
||
def test_tools_sind_lesend_und_einheitlich(service):
|
||
tools = {t.name: t for t in service.call("__list__").tools}
|
||
assert set(tools) == {"check_connection", "list_tables", "describe_table", "query"}
|
||
assert all(t.annotations.read_only_hint and not t.annotations.destructive_hint for t in tools.values())
|
||
assert set(tools["query"].input_schema["properties"]) == {"sql"}
|
||
# Zugangsdaten nie als Argument
|
||
assert all("password" not in json.dumps(t.input_schema).lower() for t in tools.values())
|
||
|
||
|
||
def test_verbindung_pruefen(service, reader, database):
|
||
result = service.call("check_connection", headers=reader)
|
||
assert not result.is_error and "PostgreSQL" in text(result) and "Schreibrechte" not in text(result)
|
||
writer = {**reader, "X-DB-User": database["writer"]}
|
||
result = service.call("check_connection", headers=writer)
|
||
assert "Schreibrechte auf 1 Objekte" in text(result) and result.structured_content["rows"][0][2] == 1
|
||
|
||
|
||
def test_tabellen_und_beschreibung(service, reader):
|
||
result = service.call("list_tables", headers=reader)
|
||
assert "public.kunden (Tabelle, ca. 1200 Zeilen): Alle Kunden" in text(result)
|
||
assert "vertrieb.top (View" in text(result)
|
||
data = result.structured_content
|
||
assert data["columns"] == ["schema", "name", "type", "rows_estimate", "comment"] and data["row_count"] == 2
|
||
result = service.call("describe_table", {"table": "kunden"}, headers=reader)
|
||
assert "- umsatz: numeric – Umsatz in Euro" in text(result) and "- name: text (nie leer)" in text(result)
|
||
assert result.structured_content["sample"]["row_count"] == 5
|
||
result = service.call("describe_table", {"table": "gibt_es_nicht"}, headers=reader)
|
||
assert result.is_error and "gibt es nicht" in text(result)
|
||
|
||
|
||
def test_abfrage_mit_tabelle_und_limit(service, reader):
|
||
result = service.call("query", {"sql": "SELECT id, name, seit, umsatz FROM kunden ORDER BY id"}, headers=reader)
|
||
data = result.structured_content
|
||
assert data["columns"] == ["id", "name", "seit", "umsatz"] and data["row_count"] == 100 and data["truncated"]
|
||
assert data["rows"][0] == [1, "Kunde 1", "2024-01-02", 10.5]
|
||
assert "| 1 | Kunde 1 | 2024-01-02 | 10.50 |" in text(result) and "gekürzt auf 100" in text(result)
|
||
result = service.call("query", {"sql": "SELECT count(*) AS n, sum(umsatz) FROM kunden"}, headers=reader)
|
||
assert result.structured_content["rows"] == [[1200, 7566300.0]]
|
||
|
||
|
||
@pytest.mark.parametrize("sql, message", [
|
||
("DELETE FROM kunden", "genau eine lesende Abfrage"),
|
||
("/* x */ UPDATE kunden SET name = 'x'", "genau eine lesende Abfrage"),
|
||
("SELECT 1; DELETE FROM kunden", "SQL-Fehler"),
|
||
("WITH x AS (DELETE FROM kunden RETURNING *) SELECT * FROM x", "SQL-Fehler"),
|
||
("SELECT pg_sleep(5)", "abgebrochen"),
|
||
])
|
||
def test_nur_lesen(service, reader, database, sql, message):
|
||
# Auch mit einem Nutzer, der schreiben dürfte: die Transaktion ist schreibgeschützt
|
||
writer = {**reader, "X-DB-User": database["writer"]}
|
||
result = service.call("query", {"sql": sql}, headers=writer)
|
||
assert result.is_error and message in text(result), text(result)
|
||
count = service.call("query", {"sql": "SELECT count(*) FROM kunden"}, headers=reader)
|
||
assert count.structured_content["rows"] == [[1200]]
|
||
|
||
|
||
def test_ohne_oder_mit_falschem_schluessel(service, reader):
|
||
for key in (None, "falscher-schluessel-0000000"):
|
||
response = httpx2.post(f"{service.url}/mcp", headers={**({"Authorization": f"Bearer {key}"} if key else {}),
|
||
"Content-Type": "application/json"}, content="{}")
|
||
assert response.status_code == 401
|
||
assert httpx2.get(f"{service.url}/live").text == "ok"
|
||
assert httpx2.get(f"{service.url}/ready").json()["driver"] == "PostgreSQL"
|
||
|
||
|
||
def test_zugangsdaten_fehlen_oder_falsch(service, reader):
|
||
result = service.call("list_tables", headers={k: v for k, v in reader.items() if k != "X-DB-User"})
|
||
assert result.is_error and "X-DB-User fehlt" in text(result)
|
||
result = service.call("list_tables", headers={**reader, "X-DB-Password": "falsch"})
|
||
assert result.is_error and "Benutzer oder Passwort falsch" in text(result)
|
||
assert "falsch" not in text(result).replace("Passwort falsch", "")
|
||
|
||
|
||
def test_nur_freigegebene_hosts():
|
||
settings = Settings(api_keys={KEY: "x"}, allowed_hosts=("db1.intern", ".prod.intern"))
|
||
base = {"x-db-database": "shop", "x-db-user": "leser", "x-db-sslmode": "require"}
|
||
assert from_headers({**base, "x-db-host": "db1.intern"}, settings, 5432).host == "db1.intern"
|
||
assert from_headers({**base, "x-db-host": "a.prod.intern"}, settings, 5432).port == 5432
|
||
for host in ("db2.intern", "evilprod.intern", "localhost", "169.254.169.254"):
|
||
with pytest.raises(Invalid):
|
||
from_headers({**base, "x-db-host": host}, settings, 5432)
|
||
with pytest.raises(Invalid, match="TLS"):
|
||
from_headers({**base, "x-db-host": "db1.intern", "x-db-sslmode": "disable"}, settings, 5432)
|
||
assert "leser" in repr(from_headers({**base, "x-db-host": "db1.intern", "x-db-password": "geheim"}, settings,
|
||
5432)) and "geheim" not in repr(
|
||
from_headers({**base, "x-db-host": "db1.intern", "x-db-password": "geheim"}, settings, 5432))
|
||
|
||
|
||
def test_sql_pruefung():
|
||
assert check_query(" -- Kommentar\nSELECT 1;") == "-- Kommentar\nSELECT 1"
|
||
with pytest.raises(Invalid):
|
||
check_query("TABLE kunden") # nur, wenn der Treiber es erlaubt
|
||
assert check_query("TABLE kunden", ("select", "table")) == "TABLE kunden"
|
||
with pytest.raises(Invalid, match="länger"):
|
||
check_query("SELECT " + "1" * 20000)
|
||
|
||
|
||
def test_audit_log_ohne_passwort(service, reader, caplog):
|
||
with caplog.at_level(logging.INFO, logger="dbmcp.audit"):
|
||
service.call("query", {"sql": "SELECT 1 AS eins"}, headers=reader)
|
||
service.call("query", {"sql": "DELETE FROM kunden"}, headers=reader)
|
||
entries = [json.loads(r.getMessage()) for r in caplog.records if r.name == "dbmcp.audit"]
|
||
ok, error = entries[-2], entries[-1]
|
||
assert ok["client"] == "bifroest" and ok["tool"] == "query" and ok["outcome"] == "ok" and ok["rows"] == 1
|
||
assert ok["db"]["user"] == reader["X-DB-User"] and ok["sql"] == "SELECT 1 AS eins"
|
||
assert error["outcome"] == "error" and "lesende Abfrage" in error["error"]
|
||
assert PASSWORD not in caplog.text
|
||
|
||
|
||
def test_einstellungen_aus_umgebung():
|
||
settings = Settings.from_env({"DBMCP_API_KEYS": "bifroest:" + "a" * 20 + ",opencode:" + "b" * 20,
|
||
"DBMCP_ALLOWED_HOSTS": "DB1.intern, .prod.intern", "DBMCP_MAX_ROWS": "50"})
|
||
assert settings.api_keys == {"a" * 20: "bifroest", "b" * 20: "opencode"}
|
||
assert settings.allowed_hosts == ("db1.intern", ".prod.intern") and settings.max_rows == 50
|
||
with pytest.raises(ValueError, match="DBMCP_API_KEYS fehlt"):
|
||
Settings.from_env({})
|
||
with pytest.raises(ValueError, match="16 Zeichen"):
|
||
Settings.from_env({"DBMCP_API_KEYS": "x:kurz"})
|
||
|
||
|
||
def test_zahlen_einheitlich():
|
||
from decimal import Decimal
|
||
|
||
from dbmcp.core.result import jsonable
|
||
assert [jsonable(Decimal(v)) for v in ("10.00", "10.50", "3")] == [10.0, 10.5, 3.0]
|
||
assert jsonable(Decimal("12345678901234567.89")) == "12345678901234567.89"
|
||
assert jsonable(Decimal("NaN")) == "NaN" and jsonable(float("inf")) == "inf"
|