mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 14:15:52 +00:00
arc53-machine
080b75c81e
Every API tool counted as holding the owner's credentials, so outside callers were refused any write on one even when it sends nothing the owner stored. Now an action counts only when its headers or query parameters carry a saved value (sealed or legacy plaintext) or the tool stores credentials; the allowlist lists just those writes.
331 lines
15 KiB
Python
331 lines
15 KiB
Python
"""Tests for creating and editing tools whose secret lives on a connection."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from contextlib import contextmanager
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
from flask import Flask
|
|
from sqlalchemy import text
|
|
|
|
import docsgpt.api.user # noqa: F401 (import order: avoids the tools/tasks cycle)
|
|
from docsgpt.connectors import service
|
|
from docsgpt.security.encryption import encrypt_json
|
|
|
|
|
|
@pytest.fixture
|
|
def app():
|
|
return Flask(__name__)
|
|
|
|
|
|
@contextmanager
|
|
def _db(conn):
|
|
@contextmanager
|
|
def _yield():
|
|
yield conn
|
|
|
|
with patch.multiple("docsgpt.api.user.tools.routes", db_session=_yield, db_readonly=_yield), \
|
|
patch.multiple("docsgpt.connectors.service", db_session=_yield, db_readonly=_yield):
|
|
yield
|
|
|
|
|
|
def _call(app, resource, body, user="alice"):
|
|
with app.test_request_context("/api/tools", method="POST", json=body):
|
|
from flask import request
|
|
|
|
request.decoded_token = {"sub": user}
|
|
return resource().post()
|
|
|
|
|
|
def _telegram(token="123456:SECRETTOKEN", **extra):
|
|
return {
|
|
"name": "telegram",
|
|
"displayName": "Telegram",
|
|
"description": "Send messages",
|
|
"config": {"token": token},
|
|
"status": True,
|
|
**extra,
|
|
}
|
|
|
|
|
|
def _connection_row(conn, connection_id):
|
|
return dict(conn.execute(
|
|
text("SELECT * FROM connector_sessions WHERE id = CAST(:i AS uuid)"), {"i": connection_id}
|
|
).one()._mapping)
|
|
|
|
|
|
class TestCreateConnectedTool:
|
|
def test_pasted_key_goes_on_a_connection_not_the_tool(self, app, pg_conn):
|
|
from docsgpt.api.user.tools.routes import CreateTool
|
|
|
|
with _db(pg_conn):
|
|
resp = _call(app, CreateTool, _telegram())
|
|
assert resp.status_code == 200
|
|
body = resp.get_json()
|
|
tool = dict(pg_conn.execute(
|
|
text("SELECT * FROM user_tools WHERE id = CAST(:i AS uuid)"), {"i": body["id"]}
|
|
).one()._mapping)
|
|
assert str(tool["connection_id"]) == body["connection_id"]
|
|
assert "SECRETTOKEN" not in str(tool["config"])
|
|
assert "encrypted_credentials" not in (tool["config"] or {})
|
|
row = _connection_row(pg_conn, body["connection_id"])
|
|
assert service.read_secrets(row)["credentials"] == {"token": "123456:SECRETTOKEN"}
|
|
|
|
def test_same_key_twice_shares_one_connection(self, app, pg_conn):
|
|
from docsgpt.api.user.tools.routes import CreateTool
|
|
|
|
with _db(pg_conn):
|
|
first = _call(app, CreateTool, _telegram()).get_json()
|
|
second = _call(app, CreateTool, _telegram()).get_json()
|
|
assert first["id"] != second["id"]
|
|
assert first["connection_id"] == second["connection_id"]
|
|
|
|
def test_existing_connection_is_used_by_id(self, app, pg_conn):
|
|
from docsgpt.api.user.tools.routes import CreateTool
|
|
|
|
with _db(pg_conn):
|
|
first = _call(app, CreateTool, _telegram()).get_json()
|
|
resp = _call(app, CreateTool, _telegram(token="", connection_id=first["connection_id"]))
|
|
assert resp.status_code == 200
|
|
assert resp.get_json()["connection_id"] == first["connection_id"]
|
|
|
|
def test_existing_connection_path_still_validates_the_config(self, app, pg_conn):
|
|
from docsgpt.api.user.tools.routes import CreateTool
|
|
|
|
with _db(pg_conn):
|
|
first = _call(app, CreateTool, _telegram()).get_json()
|
|
with patch("docsgpt.api.user.tools.routes._validate_config",
|
|
return_value={"timeout": "Timeout must be between 1 and 300"}) as validate:
|
|
resp = _call(app, CreateTool, _telegram(token="", connection_id=first["connection_id"]))
|
|
assert resp.status_code == 400
|
|
# The connection supplies the secret, so a missing key is not an error.
|
|
assert validate.call_args.kwargs["has_existing_secrets"] is True
|
|
|
|
def test_someone_elses_connection_is_not_found(self, app, pg_conn):
|
|
from docsgpt.api.user.tools.routes import CreateTool
|
|
|
|
with _db(pg_conn):
|
|
victims = _call(app, CreateTool, _telegram(), user="victim").get_json()
|
|
resp = _call(app, CreateTool, _telegram(token="", connection_id=victims["connection_id"]))
|
|
assert resp.status_code == 404
|
|
|
|
def test_connection_of_another_service_is_not_found(self, app, pg_conn):
|
|
from docsgpt.api.user.tools.routes import CreateTool
|
|
|
|
with _db(pg_conn):
|
|
telegram = _call(app, CreateTool, _telegram()).get_json()
|
|
resp = _call(app, CreateTool, {
|
|
**_telegram(token=""), "name": "ntfy", "config": {}, "connection_id": telegram["connection_id"],
|
|
})
|
|
assert resp.status_code == 404
|
|
|
|
def test_missing_key_is_a_validation_error(self, app, pg_conn):
|
|
from docsgpt.api.user.tools.routes import CreateTool
|
|
|
|
with _db(pg_conn):
|
|
resp = _call(app, CreateTool, _telegram(token=""))
|
|
assert resp.status_code == 400
|
|
|
|
def test_disabled_connector_is_refused(self, app, pg_conn):
|
|
from docsgpt.api.user.tools.routes import CreateTool
|
|
from docsgpt.storage.db.repositories.connector_policies import ConnectorPoliciesRepository
|
|
|
|
ConnectorPoliciesRepository(pg_conn).upsert("telegram", enabled=False)
|
|
with _db(pg_conn):
|
|
resp = _call(app, CreateTool, _telegram())
|
|
assert resp.status_code == 403
|
|
|
|
def test_disabled_connector_is_refused_on_the_default_key(self, app, pg_conn):
|
|
from docsgpt.api.user.tools.routes import CreateTool
|
|
from docsgpt.storage.db.repositories.connector_policies import ConnectorPoliciesRepository
|
|
|
|
ConnectorPoliciesRepository(pg_conn).upsert("telegram", enabled=False)
|
|
with _db(pg_conn), patch.object(service, "ensure_can_store_credentials",
|
|
side_effect=service.EncryptionKeyNotConfigured("set a key")):
|
|
resp = _call(app, CreateTool, _telegram())
|
|
assert resp.status_code == 403
|
|
assert pg_conn.execute(text("SELECT count(*) FROM user_tools")).scalar() == 0
|
|
|
|
def test_default_key_on_multi_user_falls_back_to_the_tool(self, app, pg_conn):
|
|
from docsgpt.api.user.tools.routes import CreateTool
|
|
|
|
with _db(pg_conn), patch.object(service, "ensure_can_store_credentials",
|
|
side_effect=service.EncryptionKeyNotConfigured("set a key")):
|
|
resp = _call(app, CreateTool, _telegram())
|
|
assert resp.status_code == 200
|
|
tool = dict(pg_conn.execute(
|
|
text("SELECT * FROM user_tools WHERE id = CAST(:i AS uuid)"), {"i": resp.get_json()["id"]}
|
|
).one()._mapping)
|
|
assert tool["connection_id"] is None
|
|
assert "encrypted_credentials" in tool["config"]
|
|
|
|
|
|
class TestUpdateConnectedTool:
|
|
def _create(self, app, pg_conn, user="alice"):
|
|
from docsgpt.api.user.tools.routes import CreateTool
|
|
|
|
with _db(pg_conn):
|
|
return _call(app, CreateTool, _telegram(), user=user).get_json()
|
|
|
|
def test_new_key_is_written_to_the_connection(self, app, pg_conn):
|
|
from docsgpt.api.user.tools.routes import UpdateTool
|
|
|
|
created = self._create(app, pg_conn)
|
|
with _db(pg_conn):
|
|
resp = _call(app, UpdateTool, {"id": created["id"], "config": {"token": "999999:ROTATED"}})
|
|
assert resp.status_code == 200
|
|
row = _connection_row(pg_conn, created["connection_id"])
|
|
assert service.read_secrets(row)["credentials"]["token"] == "999999:ROTATED"
|
|
config = pg_conn.execute(
|
|
text("SELECT config FROM user_tools WHERE id = CAST(:i AS uuid)"), {"i": created["id"]}
|
|
).scalar()
|
|
assert "ROTATED" not in str(config)
|
|
|
|
def test_new_key_replaces_unreadable_credentials(self, app, pg_conn):
|
|
"""After a lost encryption key, editing the tool is a way back in."""
|
|
from docsgpt.api.user.tools.routes import UpdateTool
|
|
|
|
created = self._create(app, pg_conn)
|
|
pg_conn.execute(
|
|
text("UPDATE connector_sessions SET encrypted_credentials = :b WHERE id = CAST(:i AS uuid)"),
|
|
{"b": encrypt_json({"credentials": {"token": "x"}}, "someone-else"), "i": created["connection_id"]},
|
|
)
|
|
with _db(pg_conn):
|
|
resp = _call(app, UpdateTool, {"id": created["id"], "config": {"token": "999999:ROTATED"}})
|
|
assert resp.status_code == 200
|
|
row = _connection_row(pg_conn, created["connection_id"])
|
|
assert service.read_secrets(row)["credentials"] == {"token": "999999:ROTATED"}
|
|
|
|
def test_editing_without_a_new_key_keeps_the_connection(self, app, pg_conn):
|
|
from docsgpt.api.user.tools.routes import UpdateTool
|
|
|
|
created = self._create(app, pg_conn)
|
|
with _db(pg_conn):
|
|
resp = _call(app, UpdateTool, {"id": created["id"], "config": {}})
|
|
assert resp.status_code == 200
|
|
row = _connection_row(pg_conn, created["connection_id"])
|
|
assert service.read_secrets(row)["credentials"]["token"] == "123456:SECRETTOKEN"
|
|
|
|
def test_rejected_edit_leaves_the_connection_untouched(self, app, pg_conn):
|
|
from docsgpt.api.user.tools.routes import UpdateTool
|
|
|
|
created = self._create(app, pg_conn)
|
|
with _db(pg_conn), patch("docsgpt.api.user.tools.routes._validate_config",
|
|
return_value={"timeout": "Timeout must be between 1 and 300"}):
|
|
resp = _call(app, UpdateTool, {"id": created["id"], "config": {"token": "999999:ROTATED"}})
|
|
assert resp.status_code == 400
|
|
row = _connection_row(pg_conn, created["connection_id"])
|
|
assert service.read_secrets(row)["credentials"]["token"] == "123456:SECRETTOKEN"
|
|
|
|
|
|
class TestAvailableTools:
|
|
def test_tools_of_a_turned_off_connector_are_not_offered(self, app, pg_conn):
|
|
from docsgpt.api.user.tools.routes import AvailableTools
|
|
from docsgpt.storage.db.repositories.connector_policies import ConnectorPoliciesRepository
|
|
|
|
ConnectorPoliciesRepository(pg_conn).upsert("telegram", enabled=False)
|
|
with _db(pg_conn), app.test_request_context("/api/available_tools"):
|
|
from flask import request
|
|
|
|
request.decoded_token = {"sub": "alice"}
|
|
names = {t["name"] for t in AvailableTools().get().get_json()["data"]}
|
|
assert "telegram" not in names
|
|
assert "brave" in names
|
|
|
|
def test_service_tools_use_the_connector_name(self, app, pg_conn):
|
|
"""One name everywhere: "Telegram", not the tool's "Telegram Bot"."""
|
|
from docsgpt.api.user.tools.routes import AvailableTools
|
|
|
|
with _db(pg_conn), app.test_request_context("/api/available_tools"):
|
|
from flask import request
|
|
|
|
request.decoded_token = {"sub": "alice"}
|
|
tools = {t["name"]: t for t in AvailableTools().get().get_json()["data"]}
|
|
assert tools["telegram"]["displayName"] == "Telegram"
|
|
assert tools["ntfy"]["displayName"] == "ntfy"
|
|
assert tools["postgres"]["displayName"] == "PostgreSQL"
|
|
|
|
|
|
def _telegram_connection(conn, user, label, name=None):
|
|
row = conn.execute(
|
|
text(
|
|
"INSERT INTO connector_sessions (user_id, provider, connector_key, auth_kind, status, account_label, "
|
|
"account_name, encrypted_credentials) VALUES (:u, 'telegram', 'telegram', 'api_key', 'connected', "
|
|
":l, :n, :e) RETURNING *"
|
|
),
|
|
{"u": user, "l": label, "n": name, "e": encrypt_json({"credentials": {"token": label}}, user)},
|
|
).one()
|
|
return dict(row._mapping)
|
|
|
|
|
|
class TestAccountNamesInToolNames:
|
|
def _listed(self, app, pg_conn, user="alice"):
|
|
from docsgpt.api.user.tools.routes import GetTools
|
|
|
|
with _db(pg_conn), app.test_request_context("/api/get_tools"):
|
|
from flask import request
|
|
|
|
request.decoded_token = {"sub": user}
|
|
tools = GetTools().get().get_json()["tools"]
|
|
return sorted(t["customName"] for t in tools if t.get("name") == "telegram")
|
|
|
|
def test_two_accounts_are_named_after_their_accounts(self, app, pg_conn):
|
|
for label, name in (("…aaaa", "Alerts bot"), ("…bbbb", None)):
|
|
service.ensure_connection_tools(pg_conn, "alice", _telegram_connection(pg_conn, "alice", label, name))
|
|
assert self._listed(app, pg_conn) == ["Telegram · Alerts bot", "Telegram · …bbbb"]
|
|
|
|
def test_one_account_keeps_the_plain_name(self, app, pg_conn):
|
|
connection = _telegram_connection(pg_conn, "alice", "…aaaa", "Alerts bot")
|
|
service.ensure_connection_tools(pg_conn, "alice", connection)
|
|
assert self._listed(app, pg_conn) == ["Telegram"]
|
|
|
|
def test_a_name_the_user_chose_is_kept(self, app, pg_conn):
|
|
for label in ("…aaaa", "…bbbb"):
|
|
service.ensure_connection_tools(pg_conn, "alice", _telegram_connection(pg_conn, "alice", label))
|
|
pg_conn.execute(text("UPDATE user_tools SET custom_name = 'Ops' WHERE name = 'telegram' "
|
|
"AND connection_id = (SELECT id FROM connector_sessions WHERE account_label = '…aaaa')"))
|
|
assert self._listed(app, pg_conn) == ["Ops", "Telegram · …bbbb"]
|
|
|
|
def test_the_drawer_names_the_tool_after_its_account_too(self, pg_conn):
|
|
named = _telegram_connection(pg_conn, "alice", "…aaaa", "Alerts bot")
|
|
for connection in (named, _telegram_connection(pg_conn, "alice", "…bbbb")):
|
|
service.ensure_connection_tools(pg_conn, "alice", connection)
|
|
detail = service.connection_detail(pg_conn, named)
|
|
assert detail["account_name"] == "Alerts bot"
|
|
assert detail["tools"][0]["display_name"] == "Telegram · Alerts bot"
|
|
|
|
|
|
class TestOwnerCredentialWrites:
|
|
"""The tool list names the writes an agent's API allowlist can cover."""
|
|
|
|
def _listed(self, app, pg_conn):
|
|
from docsgpt.api.user.tools.routes import GetTools
|
|
|
|
with _db(pg_conn), app.test_request_context("/api/get_tools"):
|
|
from flask import request
|
|
|
|
request.decoded_token = {"sub": "alice"}
|
|
tools = GetTools().get().get_json()["tools"]
|
|
return {t["name"]: t.get("owner_credential_writes") for t in tools if t.get("ownership") == "user"}
|
|
|
|
def test_writes_on_stored_credentials_are_listed(self, app, pg_conn):
|
|
from docsgpt.storage.db.repositories.user_tools import UserToolsRepository
|
|
|
|
repo = UserToolsRepository(pg_conn)
|
|
key = {"type": "object", "properties": {"X-Key": {"type": "string", "value": "", "has_value": True}}}
|
|
repo.create("alice", "api_tool", config={"actions": {
|
|
"status": {"url": "https://x.test/s", "method": "GET", "active": True, "headers": key},
|
|
"notify": {"url": "https://x.test/n", "method": "POST", "active": True, "headers": key},
|
|
# A write that sends nothing of the owner's is not listed.
|
|
"ping": {"url": "https://x.test/p", "method": "POST", "active": True},
|
|
}})
|
|
repo.create("alice", "mcp_tool", config={"server_url": "https://m.test/mcp", "auth_type": "bearer"},
|
|
actions=[{"name": "create_issue", "active": True}, {"name": "list_issues", "active": True}])
|
|
repo.create("alice", "read_webpage", actions=[{"name": "post_page", "active": True}])
|
|
listed = self._listed(app, pg_conn)
|
|
assert listed["api_tool"] == ["notify"]
|
|
assert listed["mcp_tool"] == ["create_issue"]
|
|
# No credentials: nothing of the owner's to write with.
|
|
assert listed["read_webpage"] == []
|