From 39b691de3c5acfe8bd1316263d9933e6a2c06ba2 Mon Sep 17 00:00:00 2001 From: VineFeeder Date: Wed, 15 Jul 2026 14:01:58 +0100 Subject: [PATCH] envied-working-example update --- voolts/API.py | 223 ------------------------ voolts/HTTP.py | 410 --------------------------------------------- voolts/HTTPAPI.py | 215 ------------------------ voolts/MySQL.py | 262 ----------------------------- voolts/SQLite.py | 220 ------------------------ voolts/__init__.py | 0 6 files changed, 1330 deletions(-) delete mode 100644 voolts/API.py delete mode 100644 voolts/HTTP.py delete mode 100644 voolts/HTTPAPI.py delete mode 100644 voolts/MySQL.py delete mode 100644 voolts/SQLite.py delete mode 100644 voolts/__init__.py diff --git a/voolts/API.py b/voolts/API.py deleted file mode 100644 index 75bc699..0000000 --- a/voolts/API.py +++ /dev/null @@ -1,223 +0,0 @@ -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/voolts/HTTP.py b/voolts/HTTP.py deleted file mode 100644 index cd09b8f..0000000 --- a/voolts/HTTP.py +++ /dev/null @@ -1,410 +0,0 @@ -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/voolts/HTTPAPI.py b/voolts/HTTPAPI.py deleted file mode 100644 index 380a10c..0000000 --- a/voolts/HTTPAPI.py +++ /dev/null @@ -1,215 +0,0 @@ -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/voolts/MySQL.py b/voolts/MySQL.py deleted file mode 100644 index ab5bde3..0000000 --- a/voolts/MySQL.py +++ /dev/null @@ -1,262 +0,0 @@ -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/voolts/SQLite.py b/voolts/SQLite.py deleted file mode 100644 index e3a1805..0000000 --- a/voolts/SQLite.py +++ /dev/null @@ -1,220 +0,0 @@ -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/voolts/__init__.py b/voolts/__init__.py deleted file mode 100644 index e69de29..0000000