From 39ff3903234d6848762e354753448d6b04bcf513 Mon Sep 17 00:00:00 2001 From: VineFeeder Date: Wed, 15 Jul 2026 14:17:47 +0100 Subject: [PATCH] for removal --- vaults/API.py | 223 ++++++++++++++++++++++++ vaults/HTTP.py | 410 +++++++++++++++++++++++++++++++++++++++++++++ vaults/HTTPAPI.py | 215 ++++++++++++++++++++++++ vaults/MySQL.py | 262 +++++++++++++++++++++++++++++ vaults/SQLite.py | 220 ++++++++++++++++++++++++ vaults/__init__.py | 0 6 files changed, 1330 insertions(+) create mode 100644 vaults/API.py create mode 100644 vaults/HTTP.py create mode 100644 vaults/HTTPAPI.py create mode 100644 vaults/MySQL.py create mode 100644 vaults/SQLite.py create mode 100644 vaults/__init__.py diff --git a/vaults/API.py b/vaults/API.py new file mode 100644 index 0000000..75bc699 --- /dev/null +++ b/vaults/API.py @@ -0,0 +1,223 @@ +from typing import Iterator, Optional, Union +from uuid import UUID + +from requests import Session + +from envied.core import __version__ +from envied.core.vault import Vault + + +class API(Vault): + """Key Vault using a simple RESTful HTTP API call.""" + + def __init__(self, name: str, uri: str, token: str, no_push: bool = False): + super().__init__(name, no_push) + self.uri = uri.rstrip("/") + self.session = Session() + self.session.headers.update({"User-Agent": f"envied.v{__version__}"}) + self.session.headers.update({"Authorization": f"Bearer {token}"}) + + def get_key(self, kid: Union[UUID, str], service: str) -> Optional[str]: + if isinstance(kid, UUID): + kid = kid.hex + + data = self.session.get( + url=f"{self.uri}/{service.lower()}/{kid}", headers={"Accept": "application/json"} + ).json() + + code = int(data.get("code", 0)) + message = data.get("message") + error = { + 0: None, + 1: Exceptions.AuthRejected, + 2: Exceptions.TooManyRequests, + 3: Exceptions.ServiceTagInvalid, + 4: Exceptions.KeyIdInvalid, + }.get(code, ValueError) + + if error: + raise error(f"{message} ({code})") + + content_key = data.get("content_key") + if not content_key: + return None + + if not isinstance(content_key, str): + raise ValueError(f"Expected {content_key} to be {str}, was {type(content_key)}") + + return content_key + + def get_keys(self, service: str) -> Iterator[tuple[str, str]]: + page = 1 + + while True: + data = self.session.get( + url=f"{self.uri}/{service.lower()}", + params={"page": page, "total": 10}, + headers={"Accept": "application/json"}, + ).json() + + code = int(data.get("code", 0)) + message = data.get("message") + error = { + 0: None, + 1: Exceptions.AuthRejected, + 2: Exceptions.TooManyRequests, + 3: Exceptions.PageInvalid, + 4: Exceptions.ServiceTagInvalid, + }.get(code, ValueError) + + if error: + raise error(f"{message} ({code})") + + content_keys = data.get("content_keys") + if content_keys: + if not isinstance(content_keys, dict): + raise ValueError(f"Expected {content_keys} to be {dict}, was {type(content_keys)}") + + for key_id, key in content_keys.items(): + yield key_id, key + + pages = int(data["pages"]) + if pages <= page: + break + + page += 1 + + def add_key(self, service: str, kid: Union[UUID, str], key: str) -> bool: + if isinstance(kid, UUID): + kid = kid.hex + + data = self.session.post( + url=f"{self.uri}/{service.lower()}/{kid}", json={"content_key": key}, headers={"Accept": "application/json"} + ).json() + + code = int(data.get("code", 0)) + message = data.get("message") + error = { + 0: None, + 1: Exceptions.AuthRejected, + 2: Exceptions.TooManyRequests, + 3: Exceptions.ServiceTagInvalid, + 4: Exceptions.KeyIdInvalid, + 5: Exceptions.ContentKeyInvalid, + }.get(code, ValueError) + + if error: + raise error(f"{message} ({code})") + + # the kid:key was new to the vault (optional) + added = bool(data.get("added")) + # the key for kid was changed/updated (optional) + updated = bool(data.get("updated")) + + return added or updated + + def add_keys(self, service: str, kid_keys: dict[Union[UUID, str], str]) -> int: + # Normalize keys + normalized_keys = {str(kid).replace("-", ""): key for kid, key in kid_keys.items()} + kid_list = list(normalized_keys.keys()) + + if not kid_list: + return 0 + + # Try batches starting at 500, stepping down by 100 on failure, fallback to 1 + batch_size = 500 + total_added = 0 + i = 0 + + while i < len(kid_list): + batch_kids = kid_list[i : i + batch_size] + batch_keys = {kid: normalized_keys[kid] for kid in batch_kids} + + try: + response = self.session.post( + url=f"{self.uri}/{service.lower()}", + json={"content_keys": batch_keys}, + headers={"Accept": "application/json"}, + ) + + # Check for HTTP errors that suggest batch is too large + if response.status_code in (413, 414, 400) and batch_size > 1: + if batch_size > 100: + batch_size -= 100 + else: + batch_size = 1 + continue + + data = response.json() + except Exception: + # JSON decode error or connection issue - try smaller batch + if batch_size > 1: + if batch_size > 100: + batch_size -= 100 + else: + batch_size = 1 + continue + raise + + code = int(data.get("code", 0)) + message = data.get("message") + error = { + 0: None, + 1: Exceptions.AuthRejected, + 2: Exceptions.TooManyRequests, + 3: Exceptions.ServiceTagInvalid, + 4: Exceptions.KeyIdInvalid, + 5: Exceptions.ContentKeyInvalid, + }.get(code, ValueError) + + if error: + raise error(f"{message} ({code})") + + # each kid:key that was new to the vault (optional) + added = int(data.get("added", 0)) + # each key for a kid that was changed/updated (optional) + updated = int(data.get("updated", 0)) + + total_added += added + updated + i += batch_size + + return total_added + + def get_services(self) -> Iterator[str]: + data = self.session.post(url=self.uri, headers={"Accept": "application/json"}).json() + + code = int(data.get("code", 0)) + message = data.get("message") + error = { + 0: None, + 1: Exceptions.AuthRejected, + 2: Exceptions.TooManyRequests, + }.get(code, ValueError) + + if error: + raise error(f"{message} ({code})") + + service_list = data.get("service_list", []) + + if not isinstance(service_list, list): + raise ValueError(f"Expected {service_list} to be {list}, was {type(service_list)}") + + for service in service_list: + yield service + + +class Exceptions: + class AuthRejected(Exception): + """Authentication Error Occurred, is your token valid? Do you have permission to make this call?""" + + class TooManyRequests(Exception): + """Rate Limited; Sent too many requests in a given amount of time.""" + + class PageInvalid(Exception): + """Requested page does not exist.""" + + class ServiceTagInvalid(Exception): + """The Service Tag is invalid.""" + + class KeyIdInvalid(Exception): + """The Key ID is invalid.""" + + class ContentKeyInvalid(Exception): + """The Content Key is invalid.""" diff --git a/vaults/HTTP.py b/vaults/HTTP.py new file mode 100644 index 0000000..cd09b8f --- /dev/null +++ b/vaults/HTTP.py @@ -0,0 +1,410 @@ +import json +from enum import Enum +from typing import Iterator, Optional, Union +from uuid import UUID + +from requests import Session + +from envied.core import __version__ +from envied.core.vault import Vault + + +class InsertResult(Enum): + FAILURE = 0 + SUCCESS = 1 + ALREADY_EXISTS = 2 + + +class HTTP(Vault): + """ + Key Vault using HTTP API with support for multiple API modes. + + Supported modes: + - query: Uses GET requests with query parameters + - json: Uses POST requests with JSON payloads + - decrypt_labs: Uses DecryptLabs API format (read-only) + """ + + def __init__( + self, + name: str, + host: str, + password: Optional[str] = None, + api_key: Optional[str] = None, + username: Optional[str] = None, + api_mode: str = "query", + no_push: bool = False, + ): + """ + Initialize HTTP Vault. + + Args: + name: Vault name + host: Host URL + password: Password for query mode or API token for json mode + api_key: API key (alternative to password, used for decrypt_labs mode) + username: Username (required for query mode, ignored for json/decrypt_labs mode) + api_mode: "query" for query parameters, "json" for JSON API, or "decrypt_labs" for DecryptLabs API + no_push: If True, this vault will not receive pushed keys + """ + super().__init__(name, no_push) + self.url = host + self.password = api_key or password + if not self.password: + raise ValueError("Either password or api_key is required") + + self.username = username + self.api_mode = api_mode.lower() + self.current_title = None + self.session = Session() + self.session.headers.update({"User-Agent": f"envied.v{__version__}"}) + self.api_session_id = None + + if self.api_mode == "decrypt_labs": + self.session.headers.update({"decrypt-labs-api-key": self.password}) + self.no_push = True + + # Validate configuration based on mode + if self.api_mode == "query" and not self.username: + raise ValueError("Username is required for query mode") + elif self.api_mode not in ["query", "json", "decrypt_labs"]: + raise ValueError("api_mode must be either 'query', 'json', or 'decrypt_labs'") + + def request(self, method: str, params: dict = None) -> dict: + """Make a request to the JSON API vault.""" + if self.api_mode != "json": + raise ValueError("request method is only available in json mode") + + request_payload = { + "method": method, + "params": { + **(params or {}), + "session_id": self.api_session_id, + }, + "token": self.password, + } + + r = self.session.post(self.url, json=request_payload) + + if r.status_code == 404: + return {"status": "not_found"} + + if not r.ok: + raise ValueError(f"API returned HTTP Error {r.status_code}: {r.reason.title()}") + + try: + res = r.json() + except json.JSONDecodeError: + if r.status_code == 404: + return {"status": "not_found"} + raise ValueError(f"API returned an invalid response: {r.text}") + + if res.get("status_code") != 200: + raise ValueError(f"API returned an error: {res['status_code']} - {res['message']}") + + if session_id := res.get("message", {}).get("session_id"): + self.api_session_id = session_id + + return res.get("message", res) + + def get_key(self, kid: Union[UUID, str], service: str) -> Optional[str]: + if isinstance(kid, UUID): + kid = kid.hex + + if self.api_mode == "decrypt_labs": + try: + request_payload = {"service": service.lower(), "kid": kid} + + response = self.session.post(self.url, json=request_payload) + + if not response.ok: + return None + + data = response.json() + + if data.get("message") != "success": + return None + + cached_keys = data.get("cached_keys") + if not cached_keys: + return None + + if isinstance(cached_keys, str): + try: + cached_keys = json.loads(cached_keys) + except json.JSONDecodeError: + return cached_keys + + if isinstance(cached_keys, dict): + if cached_keys.get("kid") == kid: + return cached_keys.get("key") + if kid in cached_keys: + return cached_keys[kid] + elif isinstance(cached_keys, list): + for entry in cached_keys: + if isinstance(entry, dict): + if entry.get("kid") == kid: + return entry.get("key") + elif isinstance(entry, str) and ":" in entry: + entry_kid, entry_key = entry.split(":", 1) + if entry_kid == kid: + return entry_key + + except Exception as e: + print(f"Failed to get key from DecryptLabs ({e.__class__.__name__}: {e})") + return None + return None + + elif self.api_mode == "json": + try: + params = { + "kid": kid, + "service": service.lower(), + } + + response = self.request("GetKey", params) + if response.get("status") == "not_found": + return None + keys = response.get("keys", []) + for key_entry in keys: + if isinstance(key_entry, str) and ":" in key_entry: + entry_kid, entry_key = key_entry.split(":", 1) + if entry_kid == kid: + return entry_key + elif isinstance(key_entry, dict): + if key_entry.get("kid") == kid: + return key_entry.get("key") + except Exception as e: + print(f"Failed to get key ({e.__class__.__name__}: {e})") + return None + return None + else: # query mode + response = self.session.get( + self.url, + params={"service": service.lower(), "username": self.username, "password": self.password, "kid": kid}, + ) + + data = response.json() + + if data.get("status_code") != 200 or not data.get("keys"): + return None + + return data["keys"][0]["key"] + + def get_keys(self, service: str) -> Iterator[tuple[str, str]]: + if self.api_mode == "decrypt_labs": + return iter([]) + elif self.api_mode == "json": + # JSON API doesn't support getting all keys, so return empty iterator + # This will cause the copy command to rely on the API's internal duplicate handling + return iter([]) + else: # query mode + response = self.session.get( + self.url, params={"service": service.lower(), "username": self.username, "password": self.password} + ) + + data = response.json() + + if data.get("status_code") != 200 or not data.get("keys"): + return + + for key_entry in data["keys"]: + yield key_entry["kid"], key_entry["key"] + + 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.") + + if self.api_mode == "decrypt_labs": + return False + + if isinstance(kid, UUID): + kid = kid.hex + + title = getattr(self, "current_title", None) + + if self.api_mode == "json": + try: + response = self.request( + "InsertKey", + { + "kid": kid, + "key": key, + "service": service.lower(), + "title": title, + }, + ) + if response.get("status") == "not_found": + return False + return response.get("inserted", False) + except Exception: + return False + else: # query mode + response = self.session.get( + self.url, + params={ + "service": service.lower(), + "username": self.username, + "password": self.password, + "kid": kid, + "key": key, + "title": title, + }, + ) + + data = response.json() + + return data.get("status_code") == 200 + + def add_keys(self, service: str, kid_keys: dict[Union[UUID, str], str]) -> int: + if self.api_mode == "decrypt_labs": + return 0 + + 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.") + + processed_kid_keys = { + str(kid).replace("-", "") if isinstance(kid, UUID) else kid: key for kid, key in kid_keys.items() + } + + inserted_count = 0 + title = getattr(self, "current_title", None) + + if self.api_mode == "json": + for kid, key in processed_kid_keys.items(): + try: + response = self.request( + "InsertKey", + { + "kid": kid, + "key": key, + "service": service.lower(), + "title": title, + }, + ) + if response.get("status") == "not_found": + continue + if response.get("inserted", False): + inserted_count += 1 + except Exception: + continue + else: # query mode + for kid, key in processed_kid_keys.items(): + response = self.session.get( + self.url, + params={ + "service": service.lower(), + "username": self.username, + "password": self.password, + "kid": kid, + "key": key, + "title": title, + }, + ) + + data = response.json() + + if data.get("status_code") == 200 and data.get("inserted", True): + inserted_count += 1 + + return inserted_count + + def get_services(self) -> Iterator[str]: + if self.api_mode == "decrypt_labs": + return iter([]) + elif self.api_mode == "json": + try: + response = self.request("GetServices") + services = response.get("services", []) + for service in services: + yield service + except Exception: + return iter([]) + else: # query mode + response = self.session.get( + self.url, params={"username": self.username, "password": self.password, "list_services": True} + ) + + data = response.json() + + if data.get("status_code") != 200: + return + + services = data.get("services", []) + for service in services: + yield service + + def set_title(self, title: str): + """ + Set a title to be used for the next key insertions. + This is optional and will be sent with add_key requests if available. + """ + self.current_title = title + + def insert_key_with_result( + self, service: str, kid: Union[UUID, str], key: str, title: Optional[str] = None + ) -> InsertResult: + """ + Insert a key and return detailed result information. + This method provides more granular feedback than the standard add_key method. + Available in both API modes. + """ + if not key or key.count("0") == len(key): + raise ValueError("You cannot add a NULL Content Key to a Vault.") + + if self.api_mode == "decrypt_labs": + return InsertResult.FAILURE + + if isinstance(kid, UUID): + kid = kid.hex + + if title is None: + title = getattr(self, "current_title", None) + + if self.api_mode == "json": + try: + response = self.request( + "InsertKey", + { + "kid": kid, + "key": key, + "service": service.lower(), + "title": title, + }, + ) + + if response.get("status") == "not_found": + return InsertResult.FAILURE + + if response.get("inserted", False): + return InsertResult.SUCCESS + else: + return InsertResult.ALREADY_EXISTS + + except Exception: + return InsertResult.FAILURE + else: # query mode + response = self.session.get( + self.url, + params={ + "service": service.lower(), + "username": self.username, + "password": self.password, + "kid": kid, + "key": key, + "title": title, + }, + ) + + try: + data = response.json() + if data.get("status_code") == 200: + if data.get("inserted", True): + return InsertResult.SUCCESS + else: + return InsertResult.ALREADY_EXISTS + else: + return InsertResult.FAILURE + except Exception: + return InsertResult.FAILURE diff --git a/vaults/HTTPAPI.py b/vaults/HTTPAPI.py new file mode 100644 index 0000000..380a10c --- /dev/null +++ b/vaults/HTTPAPI.py @@ -0,0 +1,215 @@ +import json +from typing import Iterator, Optional, Union +from uuid import UUID +from enum import Enum + +from requests import Session + +from envied.core import __version__ +from envied.core.vault import Vault + + +class InsertResult(Enum): + FAILURE = 0 + SUCCESS = 1 + ALREADY_EXISTS = 2 + + +class HTTPAPI(Vault): + """Key Vault using a structured HTTP API with JSON payloads and token authentication.""" + + def __init__(self, name: str, host: str, password: str): + super().__init__(name) + self.url = host + self.password = password # This is the API token + self.current_title = None + self.session = Session() + self.session.headers.update({"User-Agent": f"Devine v{__version__}"}) + self.api_session_id = None + + def request(self, method: str, params: dict = None) -> dict: + """Make a request to the HTTPAPI vault.""" + request_payload = { + "method": method, + "params": { + **(params or {}), + "session_id": self.api_session_id, + }, + "token": self.password, + } + + r = self.session.post(self.url, json=request_payload) + + if r.status_code == 404: + return {"status": "not_found"} + + if not r.ok: + raise ValueError(f"API returned HTTP Error {r.status_code}: {r.reason.title()}") + + try: + res = r.json() + except json.JSONDecodeError: + if r.status_code == 404: + return {"status": "not_found"} + raise ValueError(f"API returned an invalid response: {r.text}") + + if res.get("status_code") != 200: + raise ValueError(f"API returned an error: {res['status_code']} - {res['message']}") + + if session_id := res.get("message", {}).get("session_id"): + self.api_session_id = session_id + + return res.get("message", res) + + def get_key(self, kid: Union[UUID, str], service: str) -> Optional[str]: + if isinstance(kid, UUID): + kid = kid.hex + + try: + title = getattr(self, "current_title", None) + response = self.request( + "GetKey", + { + "kid": kid, + "service": service.lower(), + "title": title, + }, + ) + if response.get("status") == "not_found": + return None + keys = response.get("keys", []) + for key_entry in keys: + if key_entry["kid"] == kid: + return key_entry["key"] + except Exception as e: + print(f"Failed to get key ({e.__class__.__name__}: {e})") + return None + + return None + + def get_keys(self, service: str) -> Iterator[tuple[str, str]]: + """Get all keys for a service - this may need to be implemented based on your API.""" + try: + response = self.request( + "GetKeys", + { + "service": service.lower(), + }, + ) + keys = response.get("keys", []) + for key_entry in keys: + yield key_entry["kid"], key_entry["key"] + except Exception: + return iter([]) + + 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.") + + if isinstance(kid, UUID): + kid = kid.hex + + title = getattr(self, "current_title", None) + + try: + response = self.request( + "InsertKey", + { + "kid": kid, + "key": key, + "service": service.lower(), + "title": title, + }, + ) + if response.get("status") == "not_found": + return False + return response.get("inserted", False) + except Exception: + return False + + 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.") + + processed_kid_keys = { + str(kid).replace("-", "") if isinstance(kid, UUID) else kid: key for kid, key in kid_keys.items() + } + + inserted_count = 0 + title = getattr(self, "current_title", None) + + for kid, key in processed_kid_keys.items(): + try: + response = self.request( + "InsertKey", + { + "kid": kid, + "key": key, + "service": service.lower(), + "title": title, + }, + ) + if response.get("status") == "not_found": + continue + if response.get("inserted", False): + inserted_count += 1 + except Exception: + continue + + return inserted_count + + def get_services(self) -> Iterator[str]: + """Get available services - this may need to be implemented based on your API.""" + try: + response = self.request("GetServices") + services = response.get("services", []) + for service in services: + yield service + except Exception: + return iter([]) + + def set_title(self, title: str): + """ + Set a title to be used for the next key insertions. + This is optional and will be sent with add_key requests if available. + """ + self.current_title = title + + def insert_key_with_result( + self, service: str, kid: Union[UUID, str], key: str, title: Optional[str] = None + ) -> InsertResult: + """ + Insert a key and return detailed result information. + This method provides more granular feedback than the standard add_key method. + """ + if not key or key.count("0") == len(key): + raise ValueError("You cannot add a NULL Content Key to a Vault.") + + if isinstance(kid, UUID): + kid = kid.hex + + if title is None: + title = getattr(self, "current_title", None) + + try: + response = self.request( + "InsertKey", + { + "kid": kid, + "key": key, + "service": service.lower(), + "title": title, + }, + ) + + if response.get("status") == "not_found": + return InsertResult.FAILURE + + if response.get("inserted", False): + return InsertResult.SUCCESS + else: + return InsertResult.ALREADY_EXISTS + + except Exception: + return InsertResult.FAILURE diff --git a/vaults/MySQL.py b/vaults/MySQL.py new file mode 100644 index 0000000..ab5bde3 --- /dev/null +++ b/vaults/MySQL.py @@ -0,0 +1,262 @@ +import threading +from typing import Iterator, Optional, Union +from uuid import UUID + +import pymysql +from pymysql.cursors import DictCursor + +from envied.core.services import Services +from envied.core.vault import Vault + + +class MySQL(Vault): + """Key Vault using a remotely-accessed mysql database connection.""" + + def __init__(self, name: str, host: str, database: str, username: str, no_push: bool = False, **kwargs): + """ + All extra arguments provided via **kwargs will be sent to pymysql.connect. + This can be used to provide more specific connection information. + """ + super().__init__(name, no_push) + self.slug = f"{host}:{database}:{username}" + self.conn_factory = ConnectionFactory( + dict(host=host, db=database, user=username, cursorclass=DictCursor, **kwargs) + ) + + self.permissions = self.get_permissions() + if not self.has_permission("SELECT"): + raise PermissionError(f"MySQL vault {self.slug} has no SELECT permission.") + + def get_key(self, kid: Union[UUID, str], service: str) -> Optional[str]: + 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): + # no table, no keys, simple + return None + + conn = self.conn_factory.get() + cursor = conn.cursor() + + try: + cursor.execute( + # TODO: SQL injection risk + f"SELECT `kid`, `key_` FROM `{service}` WHERE `key_`!=%s", + ("0" * 32,), + ) + for row in cursor.fetchall(): + yield row["kid"], row["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.") + + if not self.has_permission("INSERT", table=service): + raise PermissionError(f"MySQL vault {self.slug} has no INSERT permission.") + + if not self.has_table(service): + try: + self.create_table(service) + except PermissionError: + return False + + if isinstance(kid, UUID): + kid = kid.hex + + conn = self.conn_factory.get() + cursor = conn.cursor() + + try: + cursor.execute( + # TODO: SQL injection risk + f"SELECT `id` FROM `{service}` WHERE `kid`=%s AND `key_`=%s", + (kid, key), + ) + if cursor.fetchone(): + # table already has this exact KID:KEY stored + return True + cursor.execute( + # TODO: SQL injection risk + f"INSERT INTO `{service}` (kid, key_) VALUES (%s, %s)", + (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.") + + if not self.has_permission("INSERT", table=service): + raise PermissionError(f"MySQL vault {self.slug} has no INSERT permission.") + + if not self.has_table(service): + try: + self.create_table(service) + except PermissionError: + return 0 + + 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()} + + if not kid_keys: + return 0 + + conn = self.conn_factory.get() + cursor = conn.cursor() + + try: + placeholders = ",".join(["%s"] * len(kid_keys)) + cursor.execute(f"SELECT kid FROM `{service}` WHERE kid IN ({placeholders})", list(kid_keys.keys())) + existing_kids = {row["kid"] for row in cursor.fetchall()} + + 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 (%s, %s)", + 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("SHOW TABLES") + for table in cursor.fetchall(): + # each entry has a key named `Tables_in_` + yield Services.get_tag(list(table.values())[0]) + finally: + cursor.close() + + def has_table(self, name: str) -> bool: + """Check if the Vault has a Table with the specified name.""" + conn = self.conn_factory.get() + cursor = conn.cursor() + + try: + cursor.execute( + "SELECT count(TABLE_NAME) FROM information_schema.TABLES WHERE TABLE_SCHEMA=%s AND TABLE_NAME=%s", + (conn.db, name), + ) + return list(cursor.fetchone().values())[0] == 1 + finally: + cursor.close() + + def create_table(self, name: str): + """Create a Table with the specified name if not yet created.""" + if self.has_table(name): + return + + if not self.has_permission("CREATE"): + raise PermissionError(f"MySQL vault {self.slug} has no CREATE permission.") + + conn = self.conn_factory.get() + cursor = conn.cursor() + + try: + cursor.execute( + # TODO: SQL injection risk + f""" + CREATE TABLE IF NOT EXISTS {name} ( + id int AUTO_INCREMENT PRIMARY KEY, + kid VARCHAR(64) NOT NULL, + key_ VARCHAR(64) NOT NULL, + UNIQUE(kid, key_) + ); + """ + ) + finally: + conn.commit() + cursor.close() + + def get_permissions(self) -> list: + """Get and parse Grants to a more easily usable list tuple array.""" + conn = self.conn_factory.get() + cursor = conn.cursor() + + try: + cursor.execute("SHOW GRANTS") + grants = cursor.fetchall() + grants = [next(iter(x.values())) for x in grants] + grants = [tuple(x[6:].split(" TO ")[0].split(" ON ")) for x in list(grants)] + grants = [ + ( + list(map(str.strip, perms.replace("ALL PRIVILEGES", "*").split(","))), + location.replace("`", "").split("."), + ) + for perms, location in grants + ] + return grants + finally: + conn.commit() + cursor.close() + + def has_permission(self, operation: str, database: Optional[str] = None, table: Optional[str] = None) -> bool: + """Check if the current connection has a specific permission.""" + grants = [x for x in self.permissions if x[0] == ["*"] or operation.upper() in x[0]] + if grants and database: + grants = [x for x in grants if x[1][0] in (database, "*")] + if grants and table: + grants = [x for x in grants if x[1][1] in (table, "*")] + return bool(grants) + + +class ConnectionFactory: + def __init__(self, con: dict): + self._con = con + self._store = threading.local() + + def _create_connection(self) -> pymysql.Connection: + return pymysql.connect(**self._con) + + def get(self) -> pymysql.Connection: + if not hasattr(self._store, "conn"): + self._store.conn = self._create_connection() + return self._store.conn diff --git a/vaults/SQLite.py b/vaults/SQLite.py new file mode 100644 index 0000000..e3a1805 --- /dev/null +++ b/vaults/SQLite.py @@ -0,0 +1,220 @@ +import sqlite3 +import threading +from pathlib import Path +from sqlite3 import Connection +from typing import Iterator, Optional, Union +from uuid import UUID + +from envied.core.services import Services +from envied.core.vault import Vault + + +class SQLite(Vault): + """Key Vault using a locally-accessed sqlite DB file.""" + + 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 + + conn = self.conn_factory.get() + cursor = conn.cursor() + + # Try both the original service name and lowercase version to handle case sensitivity issues + 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 + finally: + cursor.close() + + def get_keys(self, service: str) -> Iterator[tuple[str, str]]: + if not self.has_table(service): + # no table, no keys, simple + return None + + conn = self.conn_factory.get() + cursor = conn.cursor() + + try: + cursor.execute(f"SELECT `kid`, `key_` FROM `{service}` 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.") + + if not self.has_table(service): + self.create_table(service) + + if isinstance(kid, UUID): + kid = kid.hex + + 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(): + # table already has this exact KID:KEY stored + 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.") + + 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()} + + 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 != "sqlite_sequence": + yield Services.get_tag(name) + finally: + cursor.close() + + def has_table(self, name: str) -> bool: + """Check 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): + """Create a Table with the specified name if not yet created.""" + 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 diff --git a/vaults/__init__.py b/vaults/__init__.py new file mode 100644 index 0000000..e69de29