Script to create/delete users in Moonraker's database

moonraker-users.py

# moonraker-users.py (by @mildmojo)
#
# Script to add/delete users from Moonraker's LMDB database.
# This would've been so much less painful if they'd used SQLite.
# But good news! They're switching to SQLite soon (posted Dec 2023):
# https://github.com/Arksine/moonraker/issues/781#issuecomment-1871333293
#
# Run this script with /home/pi/moonraker-env/bin/python. For a
# Klipper deployment installed with KIAUH, something like:
#   ~/moonraker-env/bin/python moonraker-users.py --create /home/pi/printer_data/database my_user my_password
#
# Script home: https://gist.github.com/mildmojo/f5ebf9e150a70a2d11d82ac5fc9769e5

import sys
import time
import json
import hashlib
import secrets
import lmdb

USERS_TABLE = 'authorized_users'
HASH_ITER = 100000

# Loosely based on:
#   https://github.com/Arksine/moonraker/blob/db81607370021925df8fec73539fa4f2c87276b1/moonraker/components/authorization.py#L400
# Also, moonraker's DB implementation:
#   https://github.com/Arksine/moonraker/blob/db81607370021925df8fec73539fa4f2c87276b1/moonraker/components/database.py#L75

if len(sys.argv) == 1:
    print(f"Usage: {sys.argv[0]} [--create|--delete] <db-path> <username> <password>")
    sys.exit(1)

operation_arg = sys.argv[1]

if (len(sys.argv) == 1 or
   operation_arg != '--create' and operation_arg != '--delete' or
   operation_arg == '--create' and len(sys.argv) < 3 or
   operation_arg == '--delete' and len(sys.argv) < 2):
    print(f"Usage: {sys.argv[0]} [--create|--delete] <db-path> <username> <password>")
    sys.exit(1)

env = lmdb.open(path=sys.argv[2], map_size=5*1024*1024, max_dbs=2)
db = env.open_db(USERS_TABLE.encode())

username = sys.argv[3]

if operation_arg == '--create':
    password = sys.argv[4]
    source = 'moonraker'
    salt = secrets.token_bytes(32)
    hashed_pass = hashlib.pbkdf2_hmac(
        'sha256', password.encode(), salt, HASH_ITER).hex()
    user_info = {
        'username': username,
        'password': hashed_pass,
        'salt': salt.hex(),
        'source': source,
        'created_on': time.time()
    }

    with env.begin(db=db, write=True) as txn:
        if txn.get(username.encode()):
            print(f"User '{username}' already exists.")
        else:
            txn.put(username.encode(), json.dumps(user_info).encode())
            print(f"Created user '{username}'.")
elif operation_arg == '--delete':
    with env.begin(db=db, write=True) as txn:
        if txn.get(username.encode()):
            txn.delete(username.encode())
            print(f"Deleted user '{username}'.")
        else:
            print(f"User '{username}' doesn't exist.")