spotnotes.app/seed_admin.py

101 lines
3.4 KiB
Python

import os
import json
import base64
from app import app, db
from models.user import User
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
def int_to_b64url(val):
byte_len = (val.bit_length() + 7) // 8
b = val.to_bytes(byte_len, 'big')
return base64.urlsafe_b64encode(b).decode('utf-8').rstrip('=')
def seed_admin_user(email="benjimanmiller@gmail.com", password="0kl4h0m4", first_name="Benjiman", last_name="Miller"):
with app.app_context():
# 1. Generate RSA-OAEP 2048 keypair
priv = rsa.generate_private_key(public_exponent=65537, key_size=2048)
pub = priv.public_key()
priv_nums = priv.private_numbers()
pub_nums = pub.public_numbers()
pub_jwk = {
'kty': 'RSA',
'n': int_to_b64url(pub_nums.n),
'e': int_to_b64url(pub_nums.e),
'alg': 'RSA-OAEP-256',
'ext': True,
'key_ops': ['encrypt']
}
priv_jwk = {
'kty': 'RSA',
'n': int_to_b64url(pub_nums.n),
'e': int_to_b64url(pub_nums.e),
'd': int_to_b64url(priv_nums.d),
'p': int_to_b64url(priv_nums.p),
'q': int_to_b64url(priv_nums.q),
'dp': int_to_b64url(priv_nums.dmp1),
'dq': int_to_b64url(priv_nums.dmq1),
'qi': int_to_b64url(priv_nums.iqmp),
'alg': 'RSA-OAEP-256',
'ext': True,
'key_ops': ['decrypt']
}
# 2. Derive KEK via PBKDF2
salt = os.urandom(16)
kdf = PBKDF2HMAC(
algorithm=hashes.SHA256(),
length=32,
salt=salt,
iterations=100000
)
kek = kdf.derive(password.encode('utf-8'))
# 3. Encrypt private key via AES-GCM
iv = os.urandom(12)
aesgcm = AESGCM(kek)
ct = aesgcm.encrypt(iv, json.dumps(priv_jwk).encode('utf-8'), None)
enc_bundle = json.dumps({
'iv': base64.b64encode(iv).decode('utf-8'),
'ct': base64.b64encode(ct).decode('utf-8')
})
salt_b64 = base64.b64encode(salt).decode('utf-8')
pub_str = json.dumps(pub_jwk)
# 4. Check if user exists
user = User.query.filter_by(email=email).first()
if user:
print(f"Updating existing user '{email}' to admin with new keys and password...")
user.first_name = first_name
user.last_name = last_name
user.is_admin = True
user.is_banned = False
user.set_password(password)
user.public_key = pub_str
user.encrypted_private_key = enc_bundle
user.encryption_salt = salt_b64
else:
print(f"Creating new admin user '{email}'...")
user = User(
email=email,
first_name=first_name,
last_name=last_name,
is_admin=True,
is_banned=False,
public_key=pub_str,
encrypted_private_key=enc_bundle,
encryption_salt=salt_b64
)
user.set_password(password)
db.session.add(user)
db.session.commit()
print(f"Successfully seeded admin user: {email} (is_admin={user.is_admin})")
if __name__ == '__main__':
seed_admin_user()