k
91 lines
3.3 KiB
Python
91 lines
3.3 KiB
Python
"""
|
|
Re-encrypt all stored HL API keys to the current enc:v2 format (H4 fix).
|
|
|
|
Upgrades, in place:
|
|
* enc:v1 blobs (legacy unsalted-SHA256 KEK derivation)
|
|
* plaintext rows (pre-encryption era)
|
|
to enc:v2 (PBKDF2-HMAC-SHA256, per-blob salt). Rows already in enc:v2 are
|
|
skipped — the script is idempotent and safe to re-run.
|
|
|
|
KEK rotation: set OLD_ENCRYPTION_KEY in the environment to decrypt with the
|
|
old KEK while encrypting with the current ENCRYPTION_KEY.
|
|
|
|
Usage (BACK UP THE DB FIRST):
|
|
DATABASE_URL=<prod-url> python scripts/reencrypt_keys.py [--dry-run]
|
|
"""
|
|
import argparse
|
|
import asyncio
|
|
import os
|
|
import sys
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from sqlalchemy import select
|
|
|
|
from app.database import AsyncSessionLocal
|
|
from app.models import Subscription
|
|
from app.services import crypto
|
|
|
|
|
|
async def main(dry_run: bool) -> None:
|
|
old_kek = os.environ.get("OLD_ENCRYPTION_KEY")
|
|
if old_kek:
|
|
# Decrypt with the old KEK, encrypt with the new one. Swapping the
|
|
# setting around each call keeps crypto.py's caching simple.
|
|
print("KEK rotation mode: decrypting with OLD_ENCRYPTION_KEY")
|
|
|
|
upgraded = skipped = failed = 0
|
|
async with AsyncSessionLocal() as db:
|
|
rows = (await db.execute(
|
|
select(Subscription).where(Subscription.hl_api_key.is_not(None))
|
|
)).scalars().all()
|
|
print(f"{len(rows)} subscription(s) with a stored HL key")
|
|
|
|
for sub in rows:
|
|
stored = sub.hl_api_key
|
|
if crypto.is_current_format(stored) and not old_kek:
|
|
skipped += 1
|
|
continue
|
|
try:
|
|
if old_kek:
|
|
current = crypto.settings.encryption_key
|
|
crypto.settings.encryption_key = old_kek
|
|
crypto._fernet_v1 = None
|
|
crypto._fernet_v2_cache.clear()
|
|
try:
|
|
plaintext = crypto.decrypt_api_key(stored)
|
|
finally:
|
|
crypto.settings.encryption_key = current
|
|
crypto._fernet_v1 = None
|
|
crypto._fernet_v2_cache.clear()
|
|
crypto._encrypt_salt = None
|
|
else:
|
|
plaintext = crypto.decrypt_api_key(stored)
|
|
new_blob = crypto.encrypt_api_key(plaintext)
|
|
# Round-trip check before touching the row.
|
|
assert crypto.decrypt_api_key(new_blob) == plaintext
|
|
if not dry_run:
|
|
sub.hl_api_key = new_blob
|
|
upgraded += 1
|
|
print(f" {sub.wallet_address}: upgraded"
|
|
f"{' (dry-run)' if dry_run else ''}")
|
|
except Exception as exc:
|
|
failed += 1
|
|
print(f" {sub.wallet_address}: FAILED — {type(exc).__name__}: {exc}")
|
|
|
|
if not dry_run and upgraded:
|
|
await db.commit()
|
|
print("Committed.")
|
|
|
|
print(f"Done: upgraded={upgraded} skipped(already v2)={skipped} failed={failed}")
|
|
if failed:
|
|
sys.exit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--dry-run", action="store_true",
|
|
help="report what would change without writing")
|
|
args = ap.parse_args()
|
|
asyncio.run(main(args.dry_run))
|