diff --git a/py/api/admin.py b/py/api/admin.py index 64d829a..12e1162 100644 --- a/py/api/admin.py +++ b/py/api/admin.py @@ -30,6 +30,7 @@ except ImportError: IMAGE_EXTENSIONS = (".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp", ".tiff") VIDEO_EXTENSIONS = (".mp4", ".avi", ".mov", ".mkv", ".webm", ".m4v", ".wmv") SCAN_BATCH_SIZE = 50 # Files per executor call during the output scan +MAX_SAFETY_BACKUPS = 10 # '.backup_*' files kept next to the live database def _collect_media_files(output_dirs, extensions): @@ -74,6 +75,73 @@ def _sse(payload): return f"data: {json.dumps(payload)}\n\n" +def _inspect_restore_upload(path): + """Checks on a restore upload beyond PromptModel.verify_database_file. + + Opens the file read-only with ``PRAGMA trusted_schema=OFF`` (so nothing + in the schema runs while it is inspected) and refuses databases that + carry triggers or views: the prompts database never needs either, and a + crafted upload could use them to run SQL against the live data later. + + Returns: + None when acceptable, otherwise a human-readable reason (blocking). + """ + uri = Path(os.path.abspath(path)).as_uri() + "?mode=ro" + try: + conn = sqlite3.connect(uri, uri=True) + except sqlite3.Error as e: + return f"Invalid SQLite database: {e}" + try: + conn.execute("PRAGMA trusted_schema=OFF") + rows = conn.execute( + "SELECT type, name FROM sqlite_master " + "WHERE type IN ('trigger', 'view') ORDER BY type, name" + ).fetchall() + except sqlite3.Error as e: + return f"Invalid SQLite database: {e}" + finally: + conn.close() + if rows: + listed = ", ".join(f"{kind} {name}" for kind, name in rows[:5]) + return ( + "Database contains triggers or views, which a prompts database " + f"never has: {listed}" + ) + return None + + +def _prune_safety_backups(db_path, keep=None): + """Delete the oldest ``.backup_*`` files beyond ``keep`` (blocking). + + Returns the paths removed. Errors deleting a file are ignored so a + stubborn old backup never fails the restore that just succeeded. + """ + if keep is None: + keep = MAX_SAFETY_BACKUPS + db_path = os.path.abspath(db_path) + directory = os.path.dirname(db_path) + prefix = os.path.basename(db_path) + ".backup_" + try: + names = os.listdir(directory) + except OSError: + return [] + candidates = [ + os.path.join(directory, name) + for name in names + if name.startswith(prefix) and os.path.isfile(os.path.join(directory, name)) + ] + candidates.sort(key=lambda p: (os.path.getmtime(p), p)) + excess = candidates[: max(0, len(candidates) - keep)] + removed = [] + for path in excess: + try: + os.remove(path) + removed.append(path) + except OSError: + continue + return removed + + def _public_helpers(): """The package's public-path helpers (imported lazily to avoid a cycle).""" from . import _public_error, _public_path @@ -790,7 +858,7 @@ class AdminRoutesMixin: db.check_hash_duplicates() ), "statistics": lambda: { - "info": db.model.get_database_info(), + "info": self._public_database_info(db.model.get_database_info()), "message": "Database statistics retrieved", }, "prune_orphaned_prompts": lambda: count_result( @@ -803,6 +871,13 @@ class AdminRoutesMixin: ), } + def _public_database_info(self, info): + """``get_database_info()`` with the absolute database_path made public.""" + info = dict(info or {}) + if info.get("database_path"): + info["database_path"] = self._public_path(info["database_path"]) + return info + @staticmethod def _hash_duplicates_result(groups): return { @@ -832,8 +907,7 @@ class AdminRoutesMixin: fd, temp_path = tempfile.mkstemp(suffix=".db") os.close(fd) try: - loop = asyncio.get_running_loop() - ok = await loop.run_in_executor(None, model.backup_database, temp_path) + ok = await self._run_in_executor(model.backup_database, temp_path) if not ok: return web.json_response( {"success": False, "error": "Failed to create database backup"}, @@ -862,7 +936,10 @@ class AdminRoutesMixin: except Exception as e: self.logger.error(f"Backup error: {e}", exc_info=True) return web.json_response( - {"success": False, "error": f"Failed to backup database: {str(e)}"}, + { + "success": False, + "error": f"Failed to backup database: {self._public_error(e)}", + }, status=500, ) @@ -919,16 +996,23 @@ class AdminRoutesMixin: ) model = self.db.model - ok, reason = model.verify_database_file(temp_path) + ok, reason = await self._run_in_executor( + model.verify_database_file, temp_path + ) + if ok: + reason = await self._run_in_executor(_inspect_restore_upload, temp_path) + ok = reason is None if not ok: return web.json_response( {"success": False, "error": reason}, status=400 ) - loop = asyncio.get_running_loop() - backup_path = await loop.run_in_executor( - None, model.restore_from_file, temp_path + backup_path = await self._run_in_executor( + model.restore_from_file, temp_path ) + removed = await self._run_in_executor(_prune_safety_backups, model.db_path) + if removed: + self.logger.info(f"Removed {len(removed)} old safety backups") self.db = PromptDatabase(model.db_path) prompt_count = self.db.model.get_database_info().get("total_prompts", 0) @@ -940,14 +1024,19 @@ class AdminRoutesMixin: f"Found {prompt_count} prompts." ), "prompt_count": prompt_count, - "backup_created": backup_path or None, + "backup_created": ( + self._public_path(backup_path) if backup_path else None + ), } ) except Exception as e: self.logger.error(f"Restore error: {e}", exc_info=True) return web.json_response( - {"success": False, "error": f"Failed to restore database: {str(e)}"}, + { + "success": False, + "error": f"Failed to restore database: {self._public_error(e)}", + }, status=500, ) finally: diff --git a/tests/test_api_admin.py b/tests/test_api_admin.py index 4eb7c47..5e45f99 100644 --- a/tests/test_api_admin.py +++ b/tests/test_api_admin.py @@ -1026,6 +1026,14 @@ class TestMaintenance(AdminAPITestCase): self.assertIn("info", data["results"]["statistics"]) self.assertEqual(data["results"]["check_consistency"]["issues_found"], 0) + async def test_statistics_database_path_is_public(self): + resp = await self._post({"operations": ["statistics"]}) + body = await resp.text() + self.assertEqual(resp.status, 200) + self.assertNotIn(self.tmpdir, body) + info = json.loads(body)["results"]["statistics"]["info"] + self.assertEqual(info["database_path"], "prompts.db") + async def test_unknown_operation_does_nothing(self): resp = await self._post({"operations": ["reboot"]}) data = await resp.json() @@ -1094,20 +1102,27 @@ class TestBackupRestoreContract(AdminAPITestCase): self.assertEqual(resp.status, 400) self.assertIsNotNone(self.api.db.get_prompt_by_hash("h1")) + def _snapshot_upload(self, extra_sql=()): + """Bytes of a self-contained copy of the live DB (WAL rows included).""" + import sqlite3 + from contextlib import closing + + snapshot = os.path.join(self.tmpdir, "snapshot.db") + with closing(sqlite3.connect(self.api.db.model.db_path)) as src: + with closing(sqlite3.connect(snapshot)) as dst: + src.backup(dst) + for statement in extra_sql: + dst.execute(statement) + dst.commit() + with open(snapshot, "rb") as f: + return f.read() + async def test_restore_valid_db_succeeds(self): import sqlite3 from contextlib import closing self.api.db.save_prompt(text="restored", prompt_hash="h1") - # Rows may still sit in the WAL file: build a self-contained copy. - # Close both raw connections, otherwise the restore cannot replace the - # live file (Windows) or its WAL/shm sidecars (disk I/O error). - snapshot = os.path.join(self.tmpdir, "snapshot.db") - with closing(sqlite3.connect(self.api.db.model.db_path)) as src: - with closing(sqlite3.connect(snapshot)) as dst: - src.backup(dst) - with open(snapshot, "rb") as f: - upload = f.read() + upload = self._snapshot_upload() resp = await self._restore("database_file", upload) data = await resp.json() self.assertEqual(resp.status, 200, data) @@ -1118,6 +1133,117 @@ class TestBackupRestoreContract(AdminAPITestCase): conn.execute("SELECT COUNT(*) FROM prompts").fetchone()[0], 1 ) + async def test_restore_backup_created_is_a_public_path(self): + self.api.db.save_prompt(text="restored", prompt_hash="h1") + resp = await self._restore("database_file", self._snapshot_upload()) + body = await resp.text() + self.assertEqual(resp.status, 200, body) + self.assertNotIn(self.tmpdir, body) + data = json.loads(body) + self.assertTrue(data["backup_created"].startswith("prompts.db.backup_")) + + async def test_restore_rejects_database_with_a_trigger(self): + self.api.db.save_prompt(text="keep me", prompt_hash="h1") + upload = self._snapshot_upload( + [ + "CREATE TRIGGER evil AFTER INSERT ON prompts " + "BEGIN DELETE FROM prompts; END" + ] + ) + resp = await self._restore("database_file", upload) + data = await resp.json() + self.assertEqual(resp.status, 400, data) + self.assertIn("trigger", data["error"].lower()) + self.assertIsNotNone(self.api.db.get_prompt_by_hash("h1")) + + async def test_restore_rejects_database_with_a_view(self): + self.api.db.save_prompt(text="keep me", prompt_hash="h1") + upload = self._snapshot_upload(["CREATE VIEW peek AS SELECT text FROM prompts"]) + resp = await self._restore("database_file", upload) + data = await resp.json() + self.assertEqual(resp.status, 400, data) + self.assertIn("view", data["error"].lower()) + + async def test_restore_verifies_in_the_executor(self): + self.api.db.save_prompt(text="restored", prompt_hash="h1") + upload = self._snapshot_upload() + calls = [] + original = self.api._run_in_executor + + async def spy(func, *args, **kwargs): + calls.append(getattr(func, "__name__", repr(func))) + return await original(func, *args, **kwargs) + + self.api._run_in_executor = spy + resp = await self._restore("database_file", upload) + self.assertEqual(resp.status, 200) + self.assertIn("verify_database_file", calls) + self.assertIn("restore_from_file", calls) + + async def test_restore_keeps_at_most_ten_safety_backups(self): + import time + + from py.api.admin import MAX_SAFETY_BACKUPS + + self.assertEqual(MAX_SAFETY_BACKUPS, 10) + db_path = self.api.db.model.db_path + stale = [] + now = time.time() + for i in range(12): + path = f"{db_path}.backup_202001{i + 1:02d}_000000" + with open(path, "wb") as f: + f.write(b"old") + os.utime(path, (now - 10_000 + i, now - 10_000 + i)) + stale.append(path) + unrelated = os.path.join(self.tmpdir, "other.db.backup_20200101_000000") + with open(unrelated, "wb") as f: + f.write(b"x") + self.api.db.save_prompt(text="restored", prompt_hash="h1") + + resp = await self._restore("database_file", self._snapshot_upload()) + + self.assertEqual(resp.status, 200) + remaining = sorted( + name + for name in os.listdir(self.tmpdir) + if name.startswith("prompts.db.backup_") + ) + self.assertEqual(len(remaining), MAX_SAFETY_BACKUPS) + data = await resp.json() + self.assertIn(data["backup_created"], remaining) + for path in stale[:3]: + self.assertFalse(os.path.exists(path), path) + self.assertTrue(os.path.exists(unrelated)) + + async def test_restore_failure_body_has_no_absolute_path(self): + self.api.db.save_prompt(text="restored", prompt_hash="h1") + upload = self._snapshot_upload() + + def boom(src_path): + raise PermissionError( + 13, "Permission denied", os.path.join(self.tmpdir, "prompts.db") + ) + + self.api.db.model.restore_from_file = boom + resp = await self._restore("database_file", upload) + body = await resp.text() + self.assertEqual(resp.status, 500) + self.assertNotIn(self.tmpdir, body) + self.assertIn("Permission denied", json.loads(body)["error"]) + + async def test_backup_failure_body_has_no_absolute_path(self): + def boom(path): + raise PermissionError( + 13, "Permission denied", os.path.join(self.tmpdir, "x.db") + ) + + self.api.db.model.backup_database = boom + resp = await self.client.request("GET", "/prompt_manager/backup") + body = await resp.text() + self.assertEqual(resp.status, 500) + self.assertNotIn(self.tmpdir, body) + self.assertIn("Permission denied", json.loads(body)["error"]) + class TestOutputScanInternals(AdminAPITestCase): diff --git a/tests/test_api_backup_restore.py b/tests/test_api_backup_restore.py index b075ea5..02c904b 100644 --- a/tests/test_api_backup_restore.py +++ b/tests/test_api_backup_restore.py @@ -5,6 +5,7 @@ stream the upload to disk, reject anything that is not a healthy prompts database and leave the live database untouched when it does. """ +import glob import os import sqlite3 import sys @@ -19,7 +20,7 @@ from aiohttp import web from aiohttp.test_utils import AioHTTPTestCase from database.operations import PromptDatabase -from py.api import PromptManagerAPI +from py.api import PromptManagerAPI, _public_path from utils.hashing import generate_prompt_hash @@ -142,8 +143,12 @@ class TestRestoreRoute(BackupRestoreTestCase): data = await resp.json() self.assertTrue(data["success"]) self.assertEqual(data["prompt_count"], 2) - self.assertTrue(os.path.exists(data["backup_created"])) - self.assertEqual(_texts(data["backup_created"]), ["old prompt"]) + # The response names the safety backup without the server's layout + backups = glob.glob(self.db_path + ".backup_*") + self.assertEqual(len(backups), 1) + self.assertEqual(data["backup_created"], _public_path(backups[0])) + self.assertNotIn(self.tmpdir, data["backup_created"]) + self.assertEqual(_texts(backups[0]), ["old prompt"]) self.assertIsNot(self.api.db, old_db) self.assertEqual(self.api.db.model.db_path, self.db_path) diff --git a/tests/test_diagnostics_config.py b/tests/test_diagnostics_config.py index f3b6db6..c05b95b 100644 --- a/tests/test_diagnostics_config.py +++ b/tests/test_diagnostics_config.py @@ -84,9 +84,18 @@ class TestAdminEndpointsUseConfigPath(unittest.TestCase): def _make_api_stub(self, db_path): """Create a minimal object that mimics the admin mixin's self.""" + from py.api import PromptManagerAPI + stub = MagicMock() stub.db.model.db_path = db_path stub.logger = MagicMock() + stub._public_path = PromptManagerAPI._public_path + stub._public_error = PromptManagerAPI._public_error + + async def run_in_executor(func, *args, **kwargs): + return func(*args, **kwargs) + + stub._run_in_executor = run_in_executor return stub def _run_async(self, coro):