From 54adb82746158dc099143b00dfaa54f0a9198d34 Mon Sep 17 00:00:00 2001 From: imSp4rky Date: Wed, 19 Aug 2026 12:11:44 -0600 Subject: [PATCH] fix(vaults): look tables up by the name the vault stores get_services() normalised table names through Services.get_tag() and then handed the result back as an identifier, so kv copy and sync skipped any table whose name differed in case or matched an alias. Yield stored names, and resolve tables case-insensitively. --- tests/core/test_vault_table_case.py | 177 ++++++++++++++++++++++++++++ unshackle/commands/kv.py | 35 ++++-- unshackle/core/vault.py | 7 +- unshackle/vaults/MySQL.py | 83 ++++++++----- unshackle/vaults/SQLite.py | 47 ++++---- 5 files changed, 283 insertions(+), 66 deletions(-) create mode 100644 tests/core/test_vault_table_case.py diff --git a/tests/core/test_vault_table_case.py b/tests/core/test_vault_table_case.py new file mode 100644 index 00000000..21606993 --- /dev/null +++ b/tests/core/test_vault_table_case.py @@ -0,0 +1,177 @@ +"""SQLite vault table names resolve case-insensitively, and survive a `kv copy` round trip. + +A vault table is named by the service's vault tag, whose case (and sometimes whose +spelling) differs from the installed service tag. Reads and writes have to find the real +table whatever its case, `get_services()` must yield names that can be fed straight back +in, and no lookup may create a second table differing only by case. +""" + +from __future__ import annotations + +import logging +import sqlite3 +import threading +from pathlib import Path +from typing import Iterator, Optional, Union +from uuid import UUID + +import pytest + +from unshackle.core.vault import Vault +from unshackle.vaults.SQLite import SQLite + +LOG = logging.getLogger("kv-test") + +KID = "0" * 31 + "1" +KEY = "a" * 32 + + +def tables(path: Path) -> list[str]: + conn = sqlite3.connect(path) + try: + return sorted(r[0] for r in conn.execute("SELECT name FROM sqlite_master WHERE type='table'")) + finally: + conn.close() + + +class UnlistableVault(Vault): + """Stands in for the API vault, which supports neither listing nor bulk dumps.""" + + def get_key(self, kid: Union[UUID, str], service: str) -> Optional[str]: + return None + + def get_keys(self, service: str) -> Iterator[tuple[str, str]]: + raise ValueError("Bulk dump is not supported.") + + def add_key(self, service: str, kid: Union[UUID, str], key: str) -> bool: + return False + + def add_keys(self, service: str, kid_keys: dict[Union[UUID, str], str]) -> int: + return 0 + + def get_services(self) -> Iterator[str]: + raise ValueError("Listing services is not supported.") + + +@pytest.fixture +def aliased_service(monkeypatch: pytest.MonkeyPatch) -> str: + """Register an SVCD service whose vault table is named after its alias, not its tag.""" + from unshackle.core import services + + monkeypatch.setattr(services, "SERVICES", [Path("SVCD") / "__init__.py"]) + monkeypatch.setattr(services, "ALIASES", {"SVCD": ("svcdalias",)}) + return "svcdalias" + + +@pytest.fixture +def vault(tmp_path: Path) -> SQLite: + """Vault whose backing DB already holds a lowercase `svca` table with one key.""" + v = SQLite("test", tmp_path / "kv.db") + v.add_keys("svca", {KID: KEY}) + return v + + +@pytest.mark.parametrize("service", ["svca", "SVCA", "SvCa"]) +def test_reads_find_table_in_any_case(vault: SQLite, service: str) -> None: + assert vault.resolve_table(service) == "svca" + assert vault.get_key(KID, service) == KEY + assert list(vault.get_keys(service)) == [(KID, KEY)] + + +def test_add_keys_writes_into_existing_table(vault: SQLite) -> None: + kid = "0" * 31 + "2" + assert vault.add_keys("SVCA", {kid: "b" * 32}) == 1 + assert tables(vault.path) == ["sqlite_sequence", "svca"] + assert vault.get_key(kid, "svca") == "b" * 32 + + +def test_add_key_writes_into_existing_table(vault: SQLite) -> None: + kid = "0" * 31 + "3" + assert vault.add_key("SVCA", kid, "c" * 32) + assert tables(vault.path) == ["sqlite_sequence", "svca"] + assert vault.get_key(kid, "svca") == "c" * 32 + + +def test_unknown_service_is_not_an_error(vault: SQLite) -> None: + assert vault.resolve_table("SVCX") is None + assert vault.get_key(KID, "SVCX") is None + assert list(vault.get_keys("SVCX")) == [] + + +def test_new_service_creates_table_with_the_given_case(tmp_path: Path) -> None: + v = SQLite("test", tmp_path / "kv.db") + v.add_keys("SVCB", {KID: KEY}) + assert tables(v.path) == ["SVCB", "sqlite_sequence"] + + +def test_synced_lowercase_table_survives_the_service_being_added_later(tmp_path: Path) -> None: + """A `kv sync` before the service exists writes `svcc`; once SVCC is installed every + lookup uses the canonical tag and must still land on that same table.""" + v = SQLite("test", tmp_path / "kv.db") + v.add_keys("svcc", {KID: KEY}) + + kid = "0" * 31 + "2" + assert v.get_key(KID, "SVCC") == KEY + assert v.add_keys("SVCC", {kid: "b" * 32}) == 1 + assert tables(v.path) == ["sqlite_sequence", "svcc"] + assert len(list(v.get_keys("SVCC"))) == 2 + + +def test_null_keys_are_rejected_and_filtered(tmp_path: Path) -> None: + v = SQLite("test", tmp_path / "kv.db") + with pytest.raises(ValueError): + v.add_keys("SVCA", {KID: "0" * 32}) + + v.create_table("svca") + conn = sqlite3.connect(v.path) + conn.execute("INSERT INTO svca (kid, key_) VALUES (?, ?)", (KID, "0" * 32)) + conn.commit() + conn.close() + + assert v.get_key(KID, "SVCA") is None + assert list(v.get_keys("SVCA")) == [] + + +def test_get_services_yields_stored_table_names(vault: SQLite, aliased_service: str) -> None: + """Names are yielded verbatim. Normalising the `svcdalias` table to the installed SVCD + tag would hand back a name that addresses no table at all.""" + vault.add_keys(aliased_service, {KID: KEY}) + assert sorted(vault.get_services()) == ["svca", aliased_service] + + +def test_every_listed_service_can_be_copied_out_and_in(tmp_path: Path, aliased_service: str) -> None: + """Every name `get_services()` lists must be readable by `get_keys()`.""" + from unshackle.commands.kv import copy_service_data + + src = SQLite("src", tmp_path / "src.db") + src.add_keys(aliased_service, {KID: KEY}) + src.add_keys("svca", {KID: KEY}) + dst = SQLite("dst", tmp_path / "dst.db") + + added = [copy_service_data(dst, src, service, LOG) for service in src.get_services()] + + assert added == [1, 1] + assert dst.get_key(KID, aliased_service) == KEY + assert dst.get_key(KID, "svca") == KEY + + +def test_copy_skips_a_service_it_cannot_read(tmp_path: Path) -> None: + """A vault that refuses bulk dumps must not abort the copy.""" + from unshackle.commands.kv import copy_service_data + + dst = SQLite("dst", tmp_path / "dst.db") + assert copy_service_data(dst, UnlistableVault("unlistable"), "svca", LOG) == 0 + + +def test_copy_skips_a_vault_that_cannot_list_services() -> None: + with pytest.raises(ValueError): + list(UnlistableVault("unlistable").get_services()) + + +def test_lookups_work_from_another_thread(vault: SQLite) -> None: + """`kv search` probes services in a ThreadPoolExecutor; connections are thread-local.""" + found: list[Optional[str]] = [] + thread = threading.Thread(target=lambda: found.append(vault.get_key(KID, "SVCA"))) + thread.start() + thread.join() + assert found == [KEY] diff --git a/unshackle/commands/kv.py b/unshackle/commands/kv.py index 4ab3e07c..1efddff7 100644 --- a/unshackle/commands/kv.py +++ b/unshackle/commands/kv.py @@ -48,7 +48,12 @@ def process_service_keys(from_vault: Vault, service: str, log: logging.Logger) - def copy_service_data(to_vault: Vault, from_vault: Vault, service: str, log: logging.Logger) -> int: """Copy data for a single service between vaults.""" - content_keys = process_service_keys(from_vault, service, log) + try: + content_keys = process_service_keys(from_vault, service, log) + except Exception as e: + log.warning(f"{service}: Could not read from {from_vault} ({e}), skipped") + return 0 + total_count = len(content_keys) if total_count == 0: @@ -124,7 +129,7 @@ def copy( raise click.UsageError("--service and --local-only are mutually exclusive.") if service: - service = Services.get_tag(service) + service = Services.get_vault_tag(service) log.info(f"Filtering by service: {service}") installed: Optional[set[str]] = None @@ -134,11 +139,18 @@ def copy( total_added = 0 for from_vault in from_vaults: - services_to_copy = [service] if service else list(from_vault.get_services()) + if service: + services_to_copy = [service] + else: + try: + services_to_copy = list(from_vault.get_services()) + except Exception as e: + log.warning(f"{from_vault.name}: cannot list services ({e}), skipped") + continue if installed is not None: before = len(services_to_copy) - services_to_copy = [s for s in services_to_copy if s and s.upper() in installed] + services_to_copy = [s for s in services_to_copy if s and Services.get_tag(s).upper() in installed] skipped = before - len(services_to_copy) if skipped: log.info(f"{from_vault.name}: skipping {skipped} service(s) not installed locally") @@ -218,7 +230,7 @@ def add(file: Path, service: str, vaults: list[str]) -> None: raise click.ClickException("You must provide at least one Vault.") log = logging.getLogger("kv") - service = Services.get_tag(service) + service = Services.get_vault_tag(service) vaults_ = load_vaults(list(vaults)) @@ -305,7 +317,7 @@ def search(kid: str, service: Optional[str], vault_name: Optional[str]) -> None: vaults_ = load_vaults(vault_names) - service_tag = Services.get_tag(service) if service else None + service_tag = Services.get_vault_tag(service) if service else None hits: list[tuple[str, str, str]] = [] for vault in vaults_: @@ -349,13 +361,14 @@ def prepare(vaults: list[str]) -> None: vaults_ = load_vaults(vaults) for vault in vaults_: - if hasattr(vault, "has_table") and hasattr(vault, "create_table"): - for service_tag in Services.get_tags(): - if vault.has_table(service_tag): - log.info(f"{vault} already has a {service_tag} Table") + if hasattr(vault, "resolve_table") and hasattr(vault, "create_table"): + for service_tag in dict.fromkeys(Services.get_vault_tag(tag) for tag in Services.get_tags()): + existing = vault.resolve_table(service_tag) + if existing: + log.info(f"{vault} already has a {existing} Table") else: try: - vault.create_table(service_tag, commit=True) + vault.create_table(service_tag) log.info(f"{vault}: Created {service_tag} Table") except PermissionError: log.error(f"{vault} user has no create table permission, skipping...") diff --git a/unshackle/core/vault.py b/unshackle/core/vault.py index a5c8747d..e7e53424 100644 --- a/unshackle/core/vault.py +++ b/unshackle/core/vault.py @@ -43,7 +43,12 @@ class Vault(metaclass=ABCMeta): @abstractmethod def get_services(self) -> Iterator[str]: - """Get a list of Service Tags from Vault.""" + """ + Get a list of Service Tags from Vault. + + Tags come back exactly as the Vault stores them, so passing one to + get_key/get_keys/add_key/add_keys reaches the namespace it came from. + """ __all__ = ("Vault",) diff --git a/unshackle/vaults/MySQL.py b/unshackle/vaults/MySQL.py index 0603d8ae..77570d1c 100644 --- a/unshackle/vaults/MySQL.py +++ b/unshackle/vaults/MySQL.py @@ -5,7 +5,6 @@ from uuid import UUID import pymysql from pymysql.cursors import DictCursor -from unshackle.core.services import Services from unshackle.core.vault import Vault @@ -31,35 +30,8 @@ class MySQL(Vault): if isinstance(kid, UUID): kid = kid.hex - service_variants = [service] - if service != service.lower(): - service_variants.append(service.lower()) - if service != service.upper(): - service_variants.append(service.upper()) - - conn = self.conn_factory.get() - cursor = conn.cursor() - - try: - for service_name in service_variants: - if not self.has_table(service_name): - continue - - cursor.execute( - # TODO: SQL injection risk - f"SELECT `id`, `key_` FROM `{service_name}` WHERE `kid`=%s AND `key_`!=%s", - (kid, "0" * 32), - ) - cek = cursor.fetchone() - if cek: - return cek["key_"] - - return None - finally: - cursor.close() - - def get_keys(self, service: str) -> Iterator[tuple[str, str]]: - if not self.has_table(service): + table = self.resolve_table(service) + if not table: return None conn = self.conn_factory.get() @@ -68,7 +40,26 @@ class MySQL(Vault): try: cursor.execute( # TODO: SQL injection risk - f"SELECT `kid`, `key_` FROM `{service}` WHERE `key_`!=%s", + f"SELECT `id`, `key_` FROM `{table}` WHERE `kid`=%s AND `key_`!=%s", + (kid, "0" * 32), + ) + cek = cursor.fetchone() + return cek["key_"] if cek else None + finally: + cursor.close() + + def get_keys(self, service: str) -> Iterator[tuple[str, str]]: + table = self.resolve_table(service) + if not table: + return None + + conn = self.conn_factory.get() + cursor = conn.cursor() + + try: + cursor.execute( + # TODO: SQL injection risk + f"SELECT `kid`, `key_` FROM `{table}` WHERE `key_`!=%s", ("0" * 32,), ) for row in cursor.fetchall(): @@ -80,6 +71,8 @@ class MySQL(Vault): if not key or key.count("0") == len(key): raise ValueError("You cannot add a NULL Content Key to a Vault.") + service = self.resolve_table(service) or service + if not self.has_permission("INSERT", table=service): raise PermissionError(f"MySQL vault {self.slug} has no INSERT permission.") @@ -119,6 +112,8 @@ class MySQL(Vault): if not key or key.count("0") == len(key): raise ValueError("You cannot add a NULL Content Key to a Vault.") + service = self.resolve_table(service) or service + if not self.has_permission("INSERT", table=service): raise PermissionError(f"MySQL vault {self.slug} has no INSERT permission.") @@ -169,10 +164,34 @@ class MySQL(Vault): cursor.execute("SHOW TABLES") for table in cursor.fetchall(): # each entry has a key named `Tables_in_` - yield Services.get_tag(list(table.values())[0]) + yield list(table.values())[0] finally: cursor.close() + def resolve_table(self, name: str) -> Optional[str]: + """ + Get the actual Table name matching `name` case-insensitively, if it exists. + + The UPPER() comparison is explicit so the result does not depend on the server's + `lower_case_table_names` setting. Where both spellings exist, exact case wins. + """ + conn = self.conn_factory.get() + cursor = conn.cursor() + + try: + cursor.execute( + "SELECT TABLE_NAME FROM information_schema.TABLES" + " WHERE TABLE_SCHEMA=%s AND UPPER(TABLE_NAME)=UPPER(%s)", + (conn.db, name), + ) + names = [list(row.values())[0] for row in cursor.fetchall()] + finally: + cursor.close() + + if name in names: + return name + return names[0] if names else None + def has_table(self, name: str) -> bool: """Check if the Vault has a Table with the specified name.""" conn = self.conn_factory.get() diff --git a/unshackle/vaults/SQLite.py b/unshackle/vaults/SQLite.py index c0dc7823..7b03274e 100644 --- a/unshackle/vaults/SQLite.py +++ b/unshackle/vaults/SQLite.py @@ -5,7 +5,6 @@ from sqlite3 import Connection from typing import Iterator, Optional, Union from uuid import UUID -from unshackle.core.services import Services from unshackle.core.vault import Vault @@ -22,40 +21,30 @@ class SQLite(Vault): if isinstance(kid, UUID): kid = kid.hex + table = self.resolve_table(service) + if not table: + return None + conn = self.conn_factory.get() cursor = conn.cursor() - service_variants = [service] - if service != service.lower(): - service_variants.append(service.lower()) - if service != service.upper(): - service_variants.append(service.upper()) - try: - for service_name in service_variants: - if not self.has_table(service_name): - continue - - cursor.execute( - f"SELECT `id`, `key_` FROM `{service_name}` WHERE `kid`=? AND `key_`!=?", (kid, "0" * 32) - ) - cek = cursor.fetchone() - if cek: - return cek[1] - - return None + cursor.execute(f"SELECT `id`, `key_` FROM `{table}` WHERE `kid`=? AND `key_`!=?", (kid, "0" * 32)) + cek = cursor.fetchone() + return cek[1] if cek else None finally: cursor.close() def get_keys(self, service: str) -> Iterator[tuple[str, str]]: - if not self.has_table(service): + table = self.resolve_table(service) + if not table: return None conn = self.conn_factory.get() cursor = conn.cursor() try: - cursor.execute(f"SELECT `kid`, `key_` FROM `{service}` WHERE `key_`!=?", ("0" * 32,)) + cursor.execute(f"SELECT `kid`, `key_` FROM `{table}` WHERE `key_`!=?", ("0" * 32,)) for kid, key_ in cursor.fetchall(): yield kid, key_ finally: @@ -65,6 +54,7 @@ class SQLite(Vault): if not key or key.count("0") == len(key): raise ValueError("You cannot add a NULL Content Key to a Vault.") + service = self.resolve_table(service) or service if not self.has_table(service): self.create_table(service) @@ -98,6 +88,7 @@ class SQLite(Vault): if not key or key.count("0") == len(key): raise ValueError("You cannot add a NULL Content Key to a Vault.") + service = self.resolve_table(service) or service if not self.has_table(service): self.create_table(service) @@ -158,7 +149,19 @@ class SQLite(Vault): cursor.execute("SELECT name FROM sqlite_master WHERE type='table';") for (name,) in cursor.fetchall(): if name != "sqlite_sequence": - yield Services.get_tag(name) + yield name + finally: + cursor.close() + + def resolve_table(self, name: str) -> Optional[str]: + """Get the actual Table name matching `name` case-insensitively, if it exists.""" + conn = self.conn_factory.get() + cursor = conn.cursor() + + try: + cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name=? COLLATE NOCASE", (name,)) + row = cursor.fetchone() + return row[0] if row else None finally: cursor.close()