kenzuya/unshackle
1
0
Fork 1
mirror of https://github.com/unshackle-dl/unshackle.git synced 2026-03-12 01:19:02 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(kv): ✨ Enhance vault loading and key copying logic

Browse Source 
* Implemented `_load_vaults` function to load and validate vaults by name.
* Improved `_copy_service_data` to handle key copying with better logging and error handling.
* Updated `copy` command to utilize the new vault loading function and streamline the process.
* Enhanced key insertion logic in MySQL and SQLite vaults to avoid inserting existing keys.
This commit is contained in:
Andy
2025-09-03 14:50:51 +00:00
parent 86bb162868
commit 2330297ea4
3 changed files with 111 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
unshackle/vaults/MySQL.py
View File

@@ -131,16 +131,27 @@ class MySQL(Vault):
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(
# TODO: SQL injection risk
f"INSERT IGNORE INTO `{service}` (kid, key_) VALUES (%s, %s)",
kid_keys.items(),
f"INSERT INTO `{service}` (kid, key_) VALUES (%s, %s)",
new_keys.items(),
)
return cursor.rowcount
return len(new_keys)
finally:
conn.commit()
cursor.close()

19
unshackle/vaults/SQLite.py
View File

@@ -102,16 +102,27 @@ class SQLite(Vault):
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(["?"] * len(kid_keys))
cursor.execute(f"SELECT kid FROM `{service}` WHERE kid IN ({placeholders})", list(kid_keys.keys()))
existing_kids = {row[0] 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(
# TODO: SQL injection risk
f"INSERT OR IGNORE INTO `{service}` (kid, key_) VALUES (?, ?)",
kid_keys.items(),
f"INSERT INTO `{service}` (kid, key_) VALUES (?, ?)",
new_keys.items(),
)
return cursor.rowcount
return len(new_keys)
finally:
conn.commit()
cursor.close()