From 6123bba96d19ff888dcf798398e3dcc63f876cf4 Mon Sep 17 00:00:00 2001 From: John Wang Date: Fri, 9 Jun 2023 11:36:38 +0800 Subject: [PATCH] feat: add reset-encrypt-key-pair cmd for self hosted mode (#325) --- api/commands.py | 32 +++++++++++++++++++++++++++++++- 1 file changed, 31 insertions(+), 1 deletion(-) diff --git a/api/commands.py b/api/commands.py index 1112333621..10beec6c9a 100644 --- a/api/commands.py +++ b/api/commands.py @@ -3,15 +3,19 @@ import random import string import click +from flask import current_app from libs.password import password_pattern, valid_password, hash_password from libs.helper import email as email_validate from extensions.ext_database import db -from models.account import InvitationCode +from libs.rsa import generate_key_pair +from models.account import InvitationCode, Tenant from models.model import Account import secrets import base64 +from models.provider import Provider + @click.command('reset-password', help='Reset the account password.') @click.option('--email', prompt=True, help='The email address of the account whose password you need to reset') @@ -73,6 +77,31 @@ def reset_email(email, new_email, email_confirm): click.echo(click.style('Congratulations!, email has been reset.', fg='green')) +@click.command('reset-encrypt-key-pair', help='Reset the asymmetric key pair of workspace for encrypt LLM credentials. ' + 'After the reset, all LLM credentials will become invalid, ' + 'requiring re-entry.' + 'Only support SELF_HOSTED mode.') +@click.confirmation_option(prompt=click.style('Are you sure you want to reset encrypt key pair?' + ' this operation cannot be rolled back!', fg='red')) +def reset_encrypt_key_pair(): + if current_app.config['EDITION'] != 'SELF_HOSTED': + click.echo(click.style('Sorry, only support SELF_HOSTED mode.', fg='red')) + return + + tenant = db.session.query(Tenant).first() + if not tenant: + click.echo(click.style('Sorry, no workspace found. Please enter /install to initialize.', fg='red')) + return + + tenant.encrypt_public_key = generate_key_pair(tenant.id) + + db.session.query(Provider).filter(Provider.provider_type == 'custom').delete() + db.session.commit() + + click.echo(click.style('Congratulations! ' + 'the asymmetric key pair of workspace {} has been reset.'.format(tenant.id), fg='green')) + + @click.command('generate-invitation-codes', help='Generate invitation codes.') @click.option('--batch', help='The batch of invitation codes.') @click.option('--count', prompt=True, help='Invitation codes count.') @@ -134,3 +163,4 @@ def register_commands(app): app.cli.add_command(reset_password) app.cli.add_command(reset_email) app.cli.add_command(generate_invitation_codes) + app.cli.add_command(reset_encrypt_key_pair)