diff --git a/wrolpi/common.py b/wrolpi/common.py index febbd83b..5942d764 100644 --- a/wrolpi/common.py +++ b/wrolpi/common.py @@ -1102,6 +1102,7 @@ class MultiFileConfig: @dataclass class WROLPiConfigValidator: + allow_database_query: bool = None archive_destination: str = None check_for_upgrades: bool = None download_daily_limit_global: int = None @@ -1348,6 +1349,7 @@ def get_config_by_file_name(file_name: str) -> ConfigFile: class WROLPiConfig(ConfigFile): file_name = 'wrolpi.yaml' default_config = dict( + allow_database_query=False, archive_destination='archive/%(domain_tag)s/%(domain)s', download_daily_limit_global=None, download_daily_limit_per_domain=None, @@ -1600,6 +1602,14 @@ class WROLPiConfig(ConfigFile): def throttle_on_startup(self, value: bool): self.update({'throttle_on_startup': value}) + @property + def allow_database_query(self) -> bool: + return self._config['allow_database_query'] + + @allow_database_query.setter + def allow_database_query(self, value: bool): + self.update({'allow_database_query': value}) + @property def wrol_mode(self) -> bool: return self._config['wrol_mode'] diff --git a/wrolpi/db_query.py b/wrolpi/db_query.py new file mode 100644 index 00000000..bcec21bb --- /dev/null +++ b/wrolpi/db_query.py @@ -0,0 +1,229 @@ +"""Read-only SQL access to the WROLPi database, for advanced users and agents (the WROLPi MCP). + +Disabled unless `allow_database_query: true` is set in wrolpi.yaml; there is no UI for this setting. + +A query can never change the database: + * The connection is opened with `mode=ro`, so SQLite itself refuses every write. + * An authorizer permits only reads (SELECT, CTEs, functions, and introspection PRAGMAs), so a statement + cannot ATTACH another file, create TEMP objects, or change a PRAGMA. +A query cannot monopolize the device: + * A progress handler interrupts it after `QUERY_TIMEOUT` seconds (this also bounds how long it pins a + WAL snapshot). + * Rows are capped at `MAX_LIMIT`, and long values are cut to `MAX_CELL_LENGTH`. +""" +import asyncio +import contextlib +import sqlite3 +import time +from dataclasses import dataclass, field +from http import HTTPStatus +from typing import List, Optional + +from sanic import Blueprint +from sanic.request import Request +from sanic_ext import validate +from sanic_ext.extensions.openapi import openapi + +from wrolpi.api_utils import json_response +from wrolpi.common import logger, get_wrolpi_config +from wrolpi.db import get_db_file +from wrolpi.errors import DatabaseQueryDisabled, InvalidDatabaseQuery +from wrolpi.schema import JSONErrorResponse + +logger = logger.getChild(__name__) + +DEFAULT_LIMIT = 500 +MAX_LIMIT = 5000 +MAX_CELL_LENGTH = 10_000 +QUERY_TIMEOUT = 10.0 +# How many SQLite VM instructions run between timeout checks. +PROGRESS_INTERVAL = 10_000 + +ALLOWED_ACTIONS = { + sqlite3.SQLITE_SELECT, + sqlite3.SQLITE_READ, + sqlite3.SQLITE_FUNCTION, + sqlite3.SQLITE_RECURSIVE, +} +# PRAGMAs which only describe the schema. A PRAGMA with an argument (e.g. `PRAGMA x=1`) is refused. +ALLOWED_PRAGMAS = { + 'foreign_key_list', + 'index_info', + 'index_list', + 'index_xinfo', + 'table_info', + 'table_list', + 'table_xinfo', +} +DENIED_FUNCTIONS = {'load_extension'} + + +def _authorizer(action: int, arg1: Optional[str], arg2: Optional[str], _db: Optional[str], _trigger: Optional[str]): + if action == sqlite3.SQLITE_PRAGMA: + if arg1 and arg1.lower() in ALLOWED_PRAGMAS: + return sqlite3.SQLITE_OK + return sqlite3.SQLITE_DENY + if action == sqlite3.SQLITE_FUNCTION and arg2 and arg2.lower() in DENIED_FUNCTIONS: + return sqlite3.SQLITE_DENY + if action in ALLOWED_ACTIONS: + return sqlite3.SQLITE_OK + return sqlite3.SQLITE_DENY + + +@contextlib.contextmanager +def read_only_connection(timeout: float = QUERY_TIMEOUT): + """A connection to the WROLPi database which can only read, and which interrupts any statement still running + `timeout` seconds after the connection was opened.""" + conn = sqlite3.connect(f'file:{get_db_file()}?mode=ro', uri=True, timeout=5, check_same_thread=False) + try: + deadline = time.monotonic() + timeout + conn.set_progress_handler(lambda: int(time.monotonic() > deadline), PROGRESS_INTERVAL) + conn.set_authorizer(_authorizer) + yield conn + finally: + conn.close() + + +def _cell(value): + if isinstance(value, bytes): + return f'' + if isinstance(value, str) and len(value) > MAX_CELL_LENGTH: + return value[:MAX_CELL_LENGTH] + f'... ' + return value + + +def _raise_query_error(e: sqlite3.Error): + message = str(e) + if isinstance(e, sqlite3.OperationalError) and message == 'interrupted': + message = f'Query exceeded the {QUERY_TIMEOUT:g} second time limit' + elif 'not authorized' in message: + message = 'Only read-only queries (SELECT) are allowed' + raise InvalidDatabaseQuery(message) from e + + +def run_query(sql: str, limit: int = DEFAULT_LIMIT) -> dict: + """Run one read-only statement, return at most `limit` rows.""" + if not sql or not sql.strip(): + raise InvalidDatabaseQuery('Query is empty') + limit = max(1, min(int(limit or DEFAULT_LIMIT), MAX_LIMIT)) + + start = time.monotonic() + with read_only_connection() as conn: + try: + curs = conn.execute(sql) + columns = [i[0] for i in curs.description or []] + rows = curs.fetchmany(limit + 1) + except sqlite3.Error as e: + _raise_query_error(e) + except sqlite3.Warning as e: + # Python refuses more than one statement. + raise InvalidDatabaseQuery(str(e)) from e + + truncated = len(rows) > limit + rows = [[_cell(i) for i in row] for row in rows[:limit]] + return dict( + columns=columns, + rows=rows, + row_count=len(rows), + truncated=truncated, + elapsed=round(time.monotonic() - start, 3), + ) + + +def _is_fts_shadow_table(name: str, virtual_tables: List[str]) -> bool: + return any(name.startswith(f'{i}_') for i in virtual_tables) + + +def get_schema() -> List[dict]: + """Describe every table and view (but not SQLite internals or FTS5's shadow tables).""" + with read_only_connection() as conn: + objects = conn.execute( + "SELECT name, type, sql FROM sqlite_master" + " WHERE type IN ('table', 'view') AND name NOT LIKE 'sqlite_%' ORDER BY name").fetchall() + virtual_tables = [name for name, _, sql in objects if sql and sql.upper().startswith('CREATE VIRTUAL TABLE')] + + tables = [] + for name, type_, sql in objects: + if _is_fts_shadow_table(name, virtual_tables): + continue + is_virtual = name in virtual_tables + columns = [ + dict(name=col_name, type=col_type, not_null=bool(not_null), primary_key=bool(pk)) + for _, col_name, col_type, not_null, _, pk in conn.execute( + 'SELECT * FROM pragma_table_info(?)', (name,)).fetchall() + ] + indexes = [] + if type_ == 'table' and not is_virtual: + for _, index_name, unique, *_ in conn.execute( + 'SELECT * FROM pragma_index_list(?)', (name,)).fetchall(): + index_columns = [i[2] for i in conn.execute( + 'SELECT * FROM pragma_index_info(?)', (index_name,)).fetchall()] + indexes.append(dict(name=index_name, unique=bool(unique), columns=index_columns)) + table = dict( + name=name, + type='virtual' if is_virtual else type_, + columns=columns, + indexes=indexes, + ) + if is_virtual or type_ == 'view': + # The definition explains an FTS5 table's tokenizer and content table, or a view's query. + table['sql'] = sql + tables.append(table) + return tables + + +def check_database_query_allowed(): + if not get_wrolpi_config().allow_database_query: + raise DatabaseQueryDisabled() + + +db_query_bp = Blueprint('DatabaseQuery', url_prefix='/api/db') + + +@dataclass +class DatabaseQueryRequest: + sql: str + limit: Optional[int] = DEFAULT_LIMIT + + +@dataclass +class DatabaseQueryResponse: + columns: List[str] + rows: List[list] + row_count: int + truncated: bool + elapsed: float + + +@dataclass +class DatabaseSchemaResponse: + tables: List[dict] = field(default_factory=list) + + +@db_query_bp.post('/query') +@openapi.definition( + description='Run one read-only SQL statement (SELECT) against the WROLPi database. Requires' + ' `allow_database_query: true` in wrolpi.yaml.', + body=DatabaseQueryRequest, +) +@openapi.response(HTTPStatus.OK, DatabaseQueryResponse) +@openapi.response(HTTPStatus.BAD_REQUEST, JSONErrorResponse) +@openapi.response(HTTPStatus.FORBIDDEN, JSONErrorResponse) +@validate(DatabaseQueryRequest) +async def post_query(_: Request, body: DatabaseQueryRequest): + check_database_query_allowed() + result = await asyncio.to_thread(run_query, body.sql, body.limit) + return json_response(result) + + +@db_query_bp.get('/schema') +@openapi.definition( + description='Describe the tables, columns, and indexes of the WROLPi database. Requires' + ' `allow_database_query: true` in wrolpi.yaml.', +) +@openapi.response(HTTPStatus.OK, DatabaseSchemaResponse) +@openapi.response(HTTPStatus.FORBIDDEN, JSONErrorResponse) +async def get_schema_request(_: Request): + check_database_query_allowed() + tables = await asyncio.to_thread(get_schema) + return json_response(dict(tables=tables)) diff --git a/wrolpi/errors.py b/wrolpi/errors.py index 4c1ad124..8485b454 100644 --- a/wrolpi/errors.py +++ b/wrolpi/errors.py @@ -229,3 +229,15 @@ class InvalidJob(APIError): code = 'INVALID_JOB' summary = 'The Job request is invalid.' status_code = HTTPStatus.BAD_REQUEST + + +class DatabaseQueryDisabled(APIError): + code = 'DATABASE_QUERY_DISABLED' + summary = 'Database queries are disabled. Set `allow_database_query: true` in wrolpi.yaml, then import it.' + status_code = HTTPStatus.FORBIDDEN + + +class InvalidDatabaseQuery(APIError): + code = 'INVALID_DATABASE_QUERY' + summary = 'The database query is invalid, not read-only, or took too long.' + status_code = HTTPStatus.BAD_REQUEST diff --git a/wrolpi/root_api.py b/wrolpi/root_api.py index 791cc9a2..31f344c6 100644 --- a/wrolpi/root_api.py +++ b/wrolpi/root_api.py @@ -30,6 +30,7 @@ from wrolpi.common import logger, get_wrolpi_config, wrol_mode_enabled, get_medi wrol_mode_check, native_only, disable_wrol_mode, enable_wrol_mode, get_global_statistics, url_strip_host, \ set_global_log_level, get_relative_to_media_directory, search_other_estimates, set_system_timezone from wrolpi.config_api import config_bp +from wrolpi.db_query import db_query_bp from wrolpi.dates import now from wrolpi.downloader import download_manager from wrolpi.errors import WROLModeEnabled, InvalidConfig, ValidationError @@ -59,6 +60,7 @@ api_app.blueprint(archive_bp) api_app.blueprint(bookmarks_bp) api_app.blueprint(collection_bp) # Unified collection endpoints api_app.blueprint(config_bp) +api_app.blueprint(db_query_bp) api_app.blueprint(docs_bp) api_app.blueprint(files_bp) api_app.blueprint(flasher_bp) diff --git a/wrolpi/test/test_db_query.py b/wrolpi/test/test_db_query.py new file mode 100644 index 00000000..081786d2 --- /dev/null +++ b/wrolpi/test/test_db_query.py @@ -0,0 +1,191 @@ +import json +from http import HTTPStatus +from unittest import mock + +import pytest + +from wrolpi import db_query +from wrolpi.common import get_wrolpi_config +from wrolpi.db import get_db_file + + +@pytest.fixture +def allow_database_query(test_wrolpi_config): + get_wrolpi_config().allow_database_query = True + yield + get_wrolpi_config().allow_database_query = False + + +async def post_query(async_client, sql: str, **kwargs): + body = dict(sql=sql, **kwargs) + return await async_client.post('/api/db/query', content=json.dumps(body)) + + +@pytest.mark.asyncio +async def test_database_query_disabled_by_default(test_session, async_client, test_wrolpi_config): + """Both endpoints refuse until the setting is enabled in wrolpi.yaml.""" + assert get_wrolpi_config().allow_database_query is False + + request, response = await post_query(async_client, 'SELECT 1') + assert response.status_code == HTTPStatus.FORBIDDEN + assert response.json['code'] == 'DATABASE_QUERY_DISABLED' + + request, response = await async_client.get('/api/db/schema') + assert response.status_code == HTTPStatus.FORBIDDEN + assert response.json['code'] == 'DATABASE_QUERY_DISABLED' + + +@pytest.mark.asyncio +async def test_database_query_enabled_by_config_import(test_session, async_client, test_wrolpi_config): + """An advanced user enables queries by editing wrolpi.yaml and importing it; there is no UI for it.""" + config = get_wrolpi_config() + config.dump_config() + contents = test_wrolpi_config.read_text() + assert 'allow_database_query: false' in contents + test_wrolpi_config.write_text(contents.replace('allow_database_query: false', 'allow_database_query: true')) + + request, response = await async_client.post('/api/config/import', + content=json.dumps(dict(file_name='wrolpi.yaml'))) + assert response.status_code == HTTPStatus.NO_CONTENT + try: + request, response = await post_query(async_client, 'SELECT 1 AS one') + assert response.status_code == HTTPStatus.OK, response.json + assert response.json['rows'] == [[1]] + finally: + config.allow_database_query = False + + +@pytest.mark.asyncio +async def test_database_query_select(test_session, async_client, allow_database_query, tag_factory): + """A SELECT returns committed rows, with column names.""" + await tag_factory('one') + await tag_factory('two') + test_session.commit() + + request, response = await post_query(async_client, 'SELECT name FROM tag ORDER BY name') + assert response.status_code == HTTPStatus.OK, response.json + assert response.json['columns'] == ['name'] + assert response.json['rows'] == [['one'], ['two']] + assert response.json['row_count'] == 2 + assert response.json['truncated'] is False + + # CTEs (including recursive ones) are reads. + sql = 'WITH RECURSIVE n(x) AS (SELECT 1 UNION ALL SELECT x + 1 FROM n WHERE x < 3) SELECT x FROM n' + request, response = await post_query(async_client, sql) + assert response.status_code == HTTPStatus.OK, response.json + assert response.json['rows'] == [[1], [2], [3]] + + +@pytest.mark.asyncio +@pytest.mark.parametrize('sql', [ + "INSERT INTO tag (name, color) VALUES ('evil', '#000000')", + "UPDATE tag SET name = 'evil'", + 'DELETE FROM tag', + 'DROP TABLE tag', + 'CREATE TABLE evil (id INTEGER)', + 'CREATE TEMP TABLE evil (id INTEGER)', + "ATTACH DATABASE ':memory:' AS evil", + 'PRAGMA journal_mode=DELETE', + 'PRAGMA foreign_keys=OFF', + 'BEGIN', + 'VACUUM', + "SELECT load_extension('evil')", + "SELECT 1; DELETE FROM tag", +]) +async def test_database_query_refuses_writes(test_session, async_client, allow_database_query, tag_factory, sql): + """Anything other than a read is refused, and the database is unchanged.""" + await tag_factory('one') + test_session.commit() + + request, response = await post_query(async_client, sql) + assert response.status_code == HTTPStatus.BAD_REQUEST, response.json + assert response.json['code'] == 'INVALID_DATABASE_QUERY' + + request, response = await post_query(async_client, 'SELECT name FROM tag') + assert response.json['rows'] == [['one']] + + +@pytest.mark.asyncio +async def test_database_query_connection_is_read_only(test_session, allow_database_query): + """Even without the authorizer, SQLite refuses to write through the connection.""" + with db_query.read_only_connection() as conn: + conn.set_authorizer(None) + with pytest.raises(Exception, match='readonly'): + conn.execute("INSERT INTO tag (name, color) VALUES ('evil', '#000000')") + + +@pytest.mark.asyncio +async def test_database_query_introspection_pragmas(test_session, async_client, allow_database_query): + """Schema-describing PRAGMAs are allowed, both as statements and as table-valued functions.""" + request, response = await post_query(async_client, 'PRAGMA table_info(tag)') + assert response.status_code == HTTPStatus.OK, response.json + assert 'name' in [row[1] for row in response.json['rows']] + + request, response = await post_query(async_client, "SELECT name FROM pragma_table_info('tag')") + assert response.status_code == HTTPStatus.OK, response.json + assert ['name'] in response.json['rows'] + + +@pytest.mark.asyncio +async def test_database_query_errors(test_session, async_client, allow_database_query): + request, response = await post_query(async_client, 'SELECT * FROM no_such_table') + assert response.status_code == HTTPStatus.BAD_REQUEST + assert 'no such table' in response.json['error'] + + request, response = await post_query(async_client, ' ') + assert response.status_code == HTTPStatus.BAD_REQUEST + assert response.json['code'] == 'INVALID_DATABASE_QUERY' + + +@pytest.mark.asyncio +async def test_database_query_limits(test_session, async_client, allow_database_query): + """Rows are capped by `limit` (and MAX_LIMIT), long text is truncated, and blobs are summarized.""" + sql = 'WITH RECURSIVE n(x) AS (SELECT 1 UNION ALL SELECT x + 1 FROM n WHERE x < 20) SELECT x FROM n' + request, response = await post_query(async_client, sql, limit=5) + assert response.status_code == HTTPStatus.OK, response.json + assert response.json['rows'] == [[1], [2], [3], [4], [5]] + assert response.json['truncated'] is True + + with mock.patch.object(db_query, 'MAX_LIMIT', 3): + request, response = await post_query(async_client, sql, limit=1000) + assert response.json['row_count'] == 3 + assert response.json['truncated'] is True + + request, response = await post_query( + async_client, f"SELECT printf('%.*c', {db_query.MAX_CELL_LENGTH + 5}, 'a'), zeroblob(16)") + text, blob = response.json['rows'][0] + assert text.startswith('a' * db_query.MAX_CELL_LENGTH + '...') + assert text.endswith('') + assert blob == '' + + +@pytest.mark.asyncio +async def test_database_query_timeout(test_session, async_client, allow_database_query): + """A query that runs too long is interrupted.""" + sql = 'WITH RECURSIVE n(x) AS (SELECT 1 UNION ALL SELECT x + 1 FROM n) SELECT count(*) FROM n' + with mock.patch.object(db_query, 'QUERY_TIMEOUT', 0.2): + request, response = await post_query(async_client, sql) + assert response.status_code == HTTPStatus.BAD_REQUEST + assert 'time limit' in response.json['error'] + + +@pytest.mark.asyncio +async def test_database_schema(test_session, async_client, allow_database_query): + request, response = await async_client.get('/api/db/schema') + assert response.status_code == HTTPStatus.OK, response.json + tables = {i['name']: i for i in response.json['tables']} + + tag = tables['tag'] + assert tag['type'] == 'table' + columns = {i['name']: i for i in tag['columns']} + assert columns['id']['primary_key'] is True + assert 'name' in columns + assert any(i['columns'] == ['name'] for i in tag['indexes']) + + # FTS5 tables are described (with their definition), but their shadow tables are not. + virtual = [i for i in tables.values() if i['type'] == 'virtual'] + assert virtual and all('fts5' in i['sql'].lower() for i in virtual) + for table in virtual: + assert f'{table["name"]}_data' not in tables + assert not any(name.startswith('sqlite_') for name in tables) + assert get_db_file().is_file()