From 792a1a9d9158dda5d2a84ee080876c6fa9d17fa5 Mon Sep 17 00:00:00 2001
From: zhuhao <37029601+hwzhuhao@users.noreply.github.com>
Date: Tue, 23 Jul 2024 14:02:41 +0800
Subject: [PATCH] add password reset function by extending the Flask command
 (#1632)

### What problem does this PR solve?
add password reset function by extending the Flask command. #1200

### Type of change

- [x] New Feature (non-breaking change which adds functionality)
---
 api/apps/__init__.py  |  3 ++-
 api/utils/commands.py | 48 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 1 deletion(-)
 create mode 100644 api/utils/commands.py

diff --git a/api/apps/__init__.py b/api/apps/__init__.py
index 2034d08fc..0c0ace7f0 100644
--- a/api/apps/__init__.py
+++ b/api/apps/__init__.py
@@ -25,7 +25,7 @@ from flask_cors import CORS
 from api.db import StatusEnum
 from api.db.db_models import close_connection
 from api.db.services import UserService
-from api.utils import CustomJSONEncoder
+from api.utils import CustomJSONEncoder, commands
 
 from flask_session import Session
 from flask_login import LoginManager
@@ -60,6 +60,7 @@ Session(app)
 login_manager = LoginManager()
 login_manager.init_app(app)
 
+commands.register_commands(app)
 
 
 def search_pages_path(pages_dir):
diff --git a/api/utils/commands.py b/api/utils/commands.py
new file mode 100644
index 000000000..0ed184d9a
--- /dev/null
+++ b/api/utils/commands.py
@@ -0,0 +1,48 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+
+import base64
+import click
+
+from flask import Flask
+from werkzeug.security import generate_password_hash
+
+from api.db.services import UserService
+
+
+@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')
+@click.option('--new-password', prompt=True, help='the new password.')
+@click.option('--password-confirm', prompt=True, help='the new password confirm.')
+def reset_password(email, new_password, password_confirm):
+    if str(new_password).strip() != str(password_confirm).strip():
+        click.echo(click.style('sorry. The two passwords do not match.', fg='red'))
+        return
+    user = UserService.query(email=email)
+    if not user:
+        click.echo(click.style('sorry. The Email is not registered!.', fg='red'))
+        return
+    encode_password = base64.b64encode(new_password.encode('utf-8')).decode('utf-8')
+    password_hash = generate_password_hash(encode_password)
+    user_dict = {
+        'password': password_hash
+    }
+    UserService.update_user(user[0].id,user_dict)
+    click.echo(click.style('Congratulations! Password has been reset.', fg='green'))
+
+
+def register_commands(app: Flask):
+    app.cli.add_command(reset_password)
\ No newline at end of file