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()