mirror of
https://github.com/unshackle-dl/unshackle.git
synced 2026-11-14 13:05:09 +00:00
303 lines
10 KiB
Python
303 lines
10 KiB
Python
import sqlite3
|
|
import threading
|
|
from pathlib import Path
|
|
from sqlite3 import Connection
|
|
from typing import Iterator, Optional, Union
|
|
from uuid import UUID
|
|
|
|
from unshackle.core.vault import Vault
|
|
|
|
|
|
class SQLite(Vault):
|
|
"""Key Vault using a locally-accessed sqlite DB file."""
|
|
|
|
local = True
|
|
|
|
def __init__(self, name: str, path: Union[str, Path], no_push: bool = False):
|
|
super().__init__(name, no_push)
|
|
self.path = Path(path).expanduser()
|
|
# TODO: Use a DictCursor or such to get fetches as dict?
|
|
self.conn_factory = ConnectionFactory(self.path)
|
|
|
|
def get_key(self, kid: Union[UUID, str], service: str) -> Optional[str]:
|
|
if isinstance(kid, UUID):
|
|
kid = kid.hex
|
|
|
|
table = self.resolve_table(service)
|
|
if not table:
|
|
return None
|
|
|
|
self.create_bad_keys_table()
|
|
conn = self.conn_factory.get()
|
|
cursor = conn.cursor()
|
|
|
|
try:
|
|
cursor.execute(
|
|
f"SELECT `id`, `key_` FROM `{table}` WHERE `kid`=? AND `key_`!=? "
|
|
"AND `key_` NOT IN (SELECT `key_` FROM `bad_keys` WHERE `kid`=?)",
|
|
(kid, "0" * 32, kid),
|
|
)
|
|
cek = cursor.fetchone()
|
|
return cek[1] 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(f"SELECT `kid`, `key_` FROM `{table}` WHERE `key_`!=?", ("0" * 32,))
|
|
for kid, key_ in cursor.fetchall():
|
|
yield kid, key_
|
|
finally:
|
|
cursor.close()
|
|
|
|
def add_key(self, service: str, kid: Union[UUID, str], key: str) -> bool:
|
|
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)
|
|
|
|
if isinstance(kid, UUID):
|
|
kid = kid.hex
|
|
|
|
if self.is_bad_key(kid.replace("-", ""), key):
|
|
return False
|
|
|
|
conn = self.conn_factory.get()
|
|
cursor = conn.cursor()
|
|
|
|
try:
|
|
cursor.execute(
|
|
# TODO: SQL injection risk
|
|
f"SELECT `id` FROM `{service}` WHERE `kid`=? AND `key_`=?",
|
|
(kid, key),
|
|
)
|
|
if cursor.fetchone():
|
|
return True
|
|
cursor.execute(
|
|
# TODO: SQL injection risk
|
|
f"INSERT INTO `{service}` (kid, key_) VALUES (?, ?)",
|
|
(kid, key),
|
|
)
|
|
finally:
|
|
conn.commit()
|
|
cursor.close()
|
|
|
|
return True
|
|
|
|
def add_keys(self, service: str, kid_keys: dict[Union[UUID, str], str]) -> int:
|
|
for kid, key in kid_keys.items():
|
|
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)
|
|
|
|
if not isinstance(kid_keys, dict):
|
|
raise ValueError(f"The kid_keys provided is not a dictionary, {kid_keys!r}")
|
|
if not all(isinstance(kid, (str, UUID)) and isinstance(key_, str) for kid, key_ in kid_keys.items()):
|
|
raise ValueError("Expecting dict with Key of str/UUID and value of str.")
|
|
|
|
if any(isinstance(kid, UUID) for kid, key_ in kid_keys.items()):
|
|
kid_keys = {kid.hex if isinstance(kid, UUID) else kid: key_ for kid, key_ in kid_keys.items()}
|
|
|
|
self.create_bad_keys_table()
|
|
bad = set(self.conn_factory.get().execute("SELECT lower(`kid`), lower(`key_`) FROM `bad_keys`").fetchall())
|
|
kid_keys = {
|
|
kid: key_ for kid, key_ in kid_keys.items() if (kid.replace("-", "").lower(), key_.lower()) not in bad
|
|
}
|
|
|
|
if not kid_keys:
|
|
return 0
|
|
|
|
conn = self.conn_factory.get()
|
|
cursor = conn.cursor()
|
|
|
|
try:
|
|
# Query existing KIDs in batches to avoid SQLite variable limit
|
|
# Try larger batch first (newer SQLite supports 32766), fall back to 500 if needed
|
|
existing_kids: set[str] = set()
|
|
kid_list = list(kid_keys.keys())
|
|
batch_size = 32000
|
|
|
|
i = 0
|
|
while i < len(kid_list):
|
|
batch = kid_list[i : i + batch_size]
|
|
placeholders = ",".join(["?"] * len(batch))
|
|
try:
|
|
cursor.execute(f"SELECT kid FROM `{service}` WHERE kid IN ({placeholders})", batch)
|
|
existing_kids.update(row[0] for row in cursor.fetchall())
|
|
i += batch_size
|
|
except sqlite3.OperationalError as e:
|
|
if "too many SQL variables" in str(e) and batch_size > 500:
|
|
batch_size = 500
|
|
continue
|
|
raise
|
|
|
|
new_keys = {kid: key for kid, key in kid_keys.items() if kid not in existing_kids}
|
|
|
|
if not new_keys:
|
|
return 0
|
|
|
|
cursor.executemany(
|
|
f"INSERT INTO `{service}` (kid, key_) VALUES (?, ?)",
|
|
new_keys.items(),
|
|
)
|
|
return len(new_keys)
|
|
finally:
|
|
conn.commit()
|
|
cursor.close()
|
|
|
|
def get_services(self) -> Iterator[str]:
|
|
conn = self.conn_factory.get()
|
|
cursor = conn.cursor()
|
|
|
|
try:
|
|
cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
|
|
for (name,) in cursor.fetchall():
|
|
if name not in ("sqlite_sequence", "bad_keys"):
|
|
yield name
|
|
finally:
|
|
cursor.close()
|
|
|
|
def is_bad_key(self, kid: Union[UUID, str], key: str) -> bool:
|
|
if isinstance(kid, UUID):
|
|
kid = kid.hex
|
|
self.create_bad_keys_table()
|
|
cursor = self.conn_factory.get().cursor()
|
|
try:
|
|
cursor.execute("SELECT 1 FROM `bad_keys` WHERE `kid`=? AND `key_`=?", (kid, key))
|
|
return cursor.fetchone() is not None
|
|
finally:
|
|
cursor.close()
|
|
|
|
def flag_bad_key(self, service: str, kid: Union[UUID, str], key: str, source: str) -> None:
|
|
if isinstance(kid, UUID):
|
|
kid = kid.hex
|
|
self.create_bad_keys_table()
|
|
conn = self.conn_factory.get()
|
|
cursor = conn.cursor()
|
|
try:
|
|
cursor.execute(
|
|
"INSERT OR IGNORE INTO `bad_keys` (service, kid, key_, source) VALUES (?, ?, ?, ?)",
|
|
(service, kid, key, source),
|
|
)
|
|
table = None if self.no_push else self.resolve_table(service)
|
|
if table:
|
|
cursor.execute(f"DELETE FROM `{table}` WHERE `kid`=? AND `key_`=?", (kid, key))
|
|
finally:
|
|
conn.commit()
|
|
cursor.close()
|
|
|
|
def unflag_bad_key(self, kid: Union[UUID, str], key: str) -> None:
|
|
if isinstance(kid, UUID):
|
|
kid = kid.hex
|
|
self.create_bad_keys_table()
|
|
conn = self.conn_factory.get()
|
|
cursor = conn.cursor()
|
|
try:
|
|
cursor.execute("DELETE FROM `bad_keys` WHERE `kid`=? AND `key_`=?", (kid, key))
|
|
finally:
|
|
conn.commit()
|
|
cursor.close()
|
|
|
|
def create_bad_keys_table(self) -> None:
|
|
"""Make the bad_keys table if it does not exist yet. Runs the DDL once per instance."""
|
|
if getattr(self, "_bad_keys_ready", False):
|
|
return
|
|
conn = self.conn_factory.get()
|
|
cursor = conn.cursor()
|
|
try:
|
|
cursor.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS bad_keys (
|
|
"service" TEXT NOT NULL,
|
|
"kid" TEXT NOT NULL COLLATE NOCASE,
|
|
"key_" TEXT NOT NULL COLLATE NOCASE,
|
|
"source" TEXT NOT NULL,
|
|
"flagged_at" TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
|
UNIQUE("kid", "key_")
|
|
);
|
|
"""
|
|
)
|
|
finally:
|
|
conn.commit()
|
|
cursor.close()
|
|
# set only after the commit: a parallel vault thread must not query a table still being made
|
|
self._bad_keys_ready = True
|
|
|
|
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()
|
|
|
|
def has_table(self, name: str) -> bool:
|
|
"""Return True if the Vault has a Table with the specified name."""
|
|
conn = self.conn_factory.get()
|
|
cursor = conn.cursor()
|
|
|
|
try:
|
|
cursor.execute("SELECT count(name) FROM sqlite_master WHERE type='table' AND name=?", (name,))
|
|
return cursor.fetchone()[0] == 1
|
|
finally:
|
|
cursor.close()
|
|
|
|
def create_table(self, name: str):
|
|
"""Make a Table with the specified name if it does not exist yet."""
|
|
if self.has_table(name):
|
|
return
|
|
|
|
conn = self.conn_factory.get()
|
|
cursor = conn.cursor()
|
|
|
|
try:
|
|
cursor.execute(
|
|
# TODO: SQL injection risk
|
|
f"""
|
|
CREATE TABLE IF NOT EXISTS {name} (
|
|
"id" INTEGER NOT NULL UNIQUE,
|
|
"kid" TEXT NOT NULL COLLATE NOCASE,
|
|
"key_" TEXT NOT NULL COLLATE NOCASE,
|
|
PRIMARY KEY("id" AUTOINCREMENT),
|
|
UNIQUE("kid", "key_")
|
|
);
|
|
"""
|
|
)
|
|
finally:
|
|
conn.commit()
|
|
cursor.close()
|
|
|
|
|
|
class ConnectionFactory:
|
|
def __init__(self, path: Union[str, Path]):
|
|
self._path = path
|
|
self._store = threading.local()
|
|
|
|
def create_connection(self) -> Connection:
|
|
conn = sqlite3.connect(self._path, timeout=30.0)
|
|
conn.execute("PRAGMA journal_mode=WAL")
|
|
conn.execute("PRAGMA synchronous=NORMAL")
|
|
conn.execute("PRAGMA busy_timeout=30000")
|
|
return conn
|
|
|
|
def get(self) -> Connection:
|
|
if not hasattr(self._store, "conn"):
|
|
self._store.conn = self.create_connection()
|
|
return self._store.conn
|