From 02589846f54c65273a2f0d1e3c93d5cab3e3a6fe Mon Sep 17 00:00:00 2001 From: ytqh Date: Sat, 15 Mar 2025 15:57:28 +0800 Subject: [PATCH] add admin loging --- api/commands.py | 93 +++++++++++- api/configs/deploy/__init__.py | 4 +- api/configs/feature/__init__.py | 5 + api/controllers/admin/auth/login.py | 139 ++++++++++++++---- api/extensions/ext_commands.py | 2 + ...2c548baeb73f_add_phone_field_to_account.py | 33 +++++ api/models/account.py | 1 + api/services/account_service.py | 122 ++++++++++++++- 8 files changed, 362 insertions(+), 37 deletions(-) create mode 100644 api/migrations/versions/2025_03_15_1511-2c548baeb73f_add_phone_field_to_account.py diff --git a/api/commands.py b/api/commands.py index 334e7daab5..e1630c7eac 100644 --- a/api/commands.py +++ b/api/commands.py @@ -5,9 +5,6 @@ import secrets from typing import Optional import click -from flask import current_app -from werkzeug.exceptions import NotFound - from configs import dify_config from constants.languages import languages from core.rag.datasource.vdb.vector_factory import Vector @@ -16,15 +13,18 @@ from core.rag.models.document import Document from events.app_event import app_was_created from extensions.ext_database import db from extensions.ext_redis import redis_client +from flask import current_app from libs.helper import email as email_validate from libs.password import hash_password, password_pattern, valid_password from libs.rsa import generate_key_pair -from models import Tenant -from models.dataset import Dataset, DatasetCollectionBinding, DocumentSegment +from models import Account, Tenant, TenantAccountJoin, TenantAccountJoinRole +from models.dataset import Dataset, DatasetCollectionBinding from models.dataset import Document as DatasetDocument -from models.model import Account, App, AppAnnotationSetting, AppMode, Conversation, MessageAnnotation +from models.dataset import DocumentSegment +from models.model import App, AppAnnotationSetting, AppMode, Conversation, MessageAnnotation from models.provider import Provider, ProviderModel from services.account_service import RegisterService, TenantService +from werkzeug.exceptions import NotFound @click.command("reset-password", help="Reset the account password.") @@ -490,11 +490,10 @@ def add_qdrant_doc_id_index(field: str): click.echo(click.style("No dataset collection bindings found.", fg="red")) return import qdrant_client + from core.rag.datasource.vdb.qdrant.qdrant_vector import QdrantConfig from qdrant_client.http.exceptions import UnexpectedResponse from qdrant_client.http.models import PayloadSchemaType - from core.rag.datasource.vdb.qdrant.qdrant_vector import QdrantConfig - for binding in bindings: if dify_config.QDRANT_URL is None: raise ValueError("Qdrant URL is required.") @@ -649,3 +648,81 @@ where sites.id is null limit 1000""" break click.echo(click.style("Fix for missing app-related sites completed successfully!", fg="green")) + + +@click.command("create-admin-with-phone", help="Create or update an admin account with a phone number.") +@click.option("--name", prompt=True, help="Admin account name") +@click.option("--phone", prompt=True, help="Admin account phone number") +@click.option("--tenant-id", prompt=False, help="Tenant ID (optional, uses first tenant if not provided)") +def create_admin_with_phone(name: str, phone: str, tenant_id: Optional[str] = None): + """ + Create or update an admin account with a phone number. + This command will create a new account if the phone doesn't exist, + or update an existing account with the specified admin role. + """ + try: + # Check if account exists with this phone number + account = db.session.query(Account).filter(Account.phone == phone).first() + + if account: + click.echo(f"Account with phone {phone} already exists. Updating account...") + + # Update account + account.name = name + db.session.commit() + else: + click.echo(f"Creating new account with phone {phone}...") + + # Create new account with phone + account = Account( + name=name, + email=f"{phone}@qingsu.chat", + phone=phone, + interface_language=languages[0], + interface_theme="light", + status="active", + ) + + db.session.add(account) + db.session.commit() + + # Get or create tenant + tenant_id = tenant_id or dify_config.DEFAULT_TENANT_ID + tenant = db.session.query(Tenant).filter(Tenant.id == tenant_id).first() + if not tenant: + click.echo(click.style(f"Tenant with ID {tenant_id} not found.", fg="red")) + return + + # Check if account is already a member of the tenant + ta_join = ( + db.session.query(TenantAccountJoin) + .filter(TenantAccountJoin.tenant_id == tenant.id, TenantAccountJoin.account_id == account.id) + .first() + ) + + if ta_join: + # Update role to end_admin + ta_join.role = TenantAccountJoinRole.END_ADMIN.value + click.echo(f"Updated account role to {TenantAccountJoinRole.END_ADMIN.value} in tenant {tenant.name}") + else: + # Add account to tenant with end_admin role + ta_join = TenantAccountJoin( + tenant_id=tenant.id, account_id=account.id, role=TenantAccountJoinRole.END_ADMIN.value + ) + db.session.add(ta_join) + click.echo(f"Added account to tenant {tenant.name} with role {TenantAccountJoinRole.END_ADMIN.value}") + + db.session.commit() + + click.echo( + click.style( + f"Successfully {'updated' if account else 'created'} admin account with phone number.", fg="green" + ) + ) + click.echo(f"Name: {name}") + click.echo(f"Phone: {phone}") + click.echo(f"Tenant: {tenant.name} (ID: {tenant.id})") + + except Exception as e: + db.session.rollback() + click.echo(click.style(f"Error: {str(e)}", fg="red")) diff --git a/api/configs/deploy/__init__.py b/api/configs/deploy/__init__.py index 51f9af786a..e8cc61c5b5 100644 --- a/api/configs/deploy/__init__.py +++ b/api/configs/deploy/__init__.py @@ -17,8 +17,8 @@ class DeploymentConfig(BaseSettings): default=False, ) - DEBUG_EMAIL_CODE_FOR_LOGIN: str = Field( - description="Default email code for login", + DEBUG_CODE_FOR_LOGIN: str = Field( + description="Default code for login", default="111111", ) diff --git a/api/configs/feature/__init__.py b/api/configs/feature/__init__.py index bd7d3b435b..5447a81826 100644 --- a/api/configs/feature/__init__.py +++ b/api/configs/feature/__init__.py @@ -31,6 +31,11 @@ class SecurityConfig(BaseSettings): default=5, ) + PHONE_CODE_LOGIN_TOKEN_EXPIRY_MINUTES: PositiveInt = Field( + description="Duration in minutes for which a phone code login token remains valid", + default=5, + ) + LOGIN_DISABLED: bool = Field( description="Whether to disable login checks", default=False, diff --git a/api/controllers/admin/auth/login.py b/api/controllers/admin/auth/login.py index 35a10e076a..c44ec1d7ba 100644 --- a/api/controllers/admin/auth/login.py +++ b/api/controllers/admin/auth/login.py @@ -1,7 +1,16 @@ -from flask import Blueprint -from flask_restful import Api, Resource # type: ignore +from typing import cast +import flask_login # type: ignore from controllers.admin import api +from controllers.service_api_with_auth.auth.error import InvalidTokenError +from controllers.service_api_with_auth.error import AccountInFreezeError, AccountNotFound +from flask import Blueprint, request +from flask_restful import Api, Resource, reqparse # type: ignore +from libs.helper import extract_remote_ip +from models.account import Account +from services.account_service import AccountService +from services.errors.account import AccountRegisterError + class SendVerificationCodeApi(Resource): def post(self): @@ -30,16 +39,36 @@ class SendVerificationCodeApi(Resource): schema: type: object properties: - success: - type: boolean - message: + result: + type: string + data: type: string 400: description: Invalid phone number format 404: description: Phone number not registered as admin """ - pass + parser = reqparse.RequestParser() + parser.add_argument("phone", type=str, required=True, location="json") + args = parser.parse_args() + + ip_address = extract_remote_ip(request) + if AccountService.is_phone_send_ip_limit(ip_address): + return {"result": "fail", "data": "Too many requests from this IP address"}, 429 + + try: + # find account by phone number & chech role is end_admin + account = AccountService.get_admin_through_phone(args["phone"]) + except AccountRegisterError: + raise AccountInFreezeError() + + if account is None: + return {"result": "fail", "data": "Phone number not registered as admin"}, 404 + + token = AccountService.send_phone_code_login(phone=args["phone"]) + + return {"result": "success", "data": token} + class LoginApi(Resource): def post(self): @@ -58,6 +87,7 @@ class LoginApi(Resource): required: - phone - code + - token properties: phone: type: string @@ -66,34 +96,76 @@ class LoginApi(Resource): code: type: string description: Verification code - example: "123456" + example: "111111" + token: + type: string + description: Verification token responses: 200: description: Login successful schema: type: object properties: - token: + result: type: string - description: JWT access token - user: + data: type: object properties: - id: - type: string - phone: - type: string - name: - type: string - role: + token: type: string - enum: [admin, super_admin] + user: + type: object + properties: + id: + type: string + phone: + type: string + name: + type: string + role: + type: string + enum: [admin, super_admin] 400: description: Invalid or expired verification code 404: description: Phone number not registered """ - pass + parser = reqparse.RequestParser() + parser.add_argument("phone", type=str, required=True, location="json") + parser.add_argument("code", type=str, required=True, location="json") + parser.add_argument("token", type=str, required=True, location="json") + args = parser.parse_args() + + # Verify the token and code + token_data = AccountService.get_phone_code_login_data(args["token"]) + if token_data is None: + return {"result": "fail", "data": "Invalid or expired token"}, 400 + + if token_data["phone"] != args["phone"]: + return {"result": "fail", "data": "Phone number does not match"}, 400 + + if token_data["code"] != args["code"]: + return {"result": "fail", "data": "Invalid verification code"}, 400 + + # Revoke the token after successful verification + AccountService.revoke_phone_code_login_token(args["token"]) + + try: + account = AccountService.get_admin_through_phone(args["phone"]) + except AccountRegisterError: + raise AccountInFreezeError() + + if account is None: + return {"result": "fail", "data": "Phone number not registered as admin"}, 404 + + # Generate token for the authenticated admin + token_pair = AccountService.login(account, ip_address=extract_remote_ip(request)) + AccountService.reset_login_error_rate_limit(args["phone"]) + + response_data = token_pair.model_dump() + + return {"result": "success", "data": response_data} + class LogoutApi(Resource): def post(self): @@ -111,12 +183,21 @@ class LogoutApi(Resource): schema: type: object properties: - success: - type: boolean + result: + type: string 401: description: Missing or invalid token """ - pass + account = cast(Account, flask_login.current_user) + + if isinstance(account, flask_login.AnonymousUserMixin): + return {"result": "success"} + + AccountService.logout(account=account) + flask_login.logout_user() + + return {"result": "success"} + class RefreshTokenApi(Resource): def post(self): @@ -148,17 +229,25 @@ class RefreshTokenApi(Resource): properties: result: type: string - example: "success" data: type: object description: New token pair data 401: description: Unauthorized, invalid or missing token """ - pass + parser = reqparse.RequestParser() + parser.add_argument("refresh_token", type=str, required=True, location="json") + args = parser.parse_args() + + try: + new_token_pair = AccountService.refresh_token(args["refresh_token"]) + return {"result": "success", "data": new_token_pair.model_dump()} + except Exception as e: + return {"result": "fail", "data": str(e)}, 401 + # Register the resources api.add_resource(SendVerificationCodeApi, '/auth/send-code') api.add_resource(LoginApi, '/auth/login') api.add_resource(LogoutApi, '/auth/logout') -api.add_resource(RefreshTokenApi, '/auth/refresh-token') \ No newline at end of file +api.add_resource(RefreshTokenApi, '/auth/refresh-token') diff --git a/api/extensions/ext_commands.py b/api/extensions/ext_commands.py index ccf0d316ca..20ec826124 100644 --- a/api/extensions/ext_commands.py +++ b/api/extensions/ext_commands.py @@ -5,6 +5,7 @@ def init_app(app: DifyApp): from commands import ( add_qdrant_doc_id_index, convert_to_agent_apps, + create_admin_with_phone, create_tenant, fix_app_site_missing, reset_email, @@ -24,6 +25,7 @@ def init_app(app: DifyApp): create_tenant, upgrade_db, fix_app_site_missing, + create_admin_with_phone, ] for cmd in cmds_to_register: app.cli.add_command(cmd) diff --git a/api/migrations/versions/2025_03_15_1511-2c548baeb73f_add_phone_field_to_account.py b/api/migrations/versions/2025_03_15_1511-2c548baeb73f_add_phone_field_to_account.py new file mode 100644 index 0000000000..0cac25651a --- /dev/null +++ b/api/migrations/versions/2025_03_15_1511-2c548baeb73f_add_phone_field_to_account.py @@ -0,0 +1,33 @@ +"""add phone field to account + +Revision ID: 2c548baeb73f +Revises: 0dab3b2ce369 +Create Date: 2025-03-15 15:11:06.752463 + +""" +from alembic import op +import models as models +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = '2c548baeb73f' +down_revision = '0dab3b2ce369' +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table('accounts', schema=None) as batch_op: + batch_op.add_column(sa.Column('phone', sa.String(length=255), nullable=True)) + + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table('accounts', schema=None) as batch_op: + batch_op.drop_column('phone') + + # ### end Alembic commands ### diff --git a/api/models/account.py b/api/models/account.py index 6f0f2b2dff..715a2cf9a4 100644 --- a/api/models/account.py +++ b/api/models/account.py @@ -24,6 +24,7 @@ class Account(UserMixin, db.Model): # type: ignore[name-defined] id: Mapped[str] = mapped_column(StringUUID, server_default=db.text("uuid_generate_v4()")) name = db.Column(db.String(255), nullable=False) email = db.Column(db.String(255), nullable=False) + phone = db.Column(db.String(255), nullable=True) password = db.Column(db.String(255), nullable=True) password_salt = db.Column(db.String(255), nullable=True) avatar = db.Column(db.String(255)) diff --git a/api/services/account_service.py b/api/services/account_service.py index 7d2490320d..54cd6e2d24 100644 --- a/api/services/account_service.py +++ b/api/services/account_service.py @@ -73,6 +73,9 @@ class AccountService: email_code_login_rate_limiter = RateLimiter( prefix="email_code_login_rate_limit", max_attempts=1, time_window=60 * 1 ) + phone_code_login_rate_limiter = RateLimiter( + prefix="phone_code_login_rate_limit", max_attempts=1, time_window=60 * 1 + ) email_code_account_deletion_rate_limiter = RateLimiter( prefix="email_code_account_deletion_rate_limit", max_attempts=1, @@ -481,7 +484,7 @@ class AccountService: raise EmailCodeLoginRateLimitExceededError() if DeploymentConfig().DEBUG: - code = dify_config.DEBUG_EMAIL_CODE_FOR_LOGIN + code = dify_config.DEBUG_CODE_FOR_LOGIN else: code = "".join([str(random.randint(0, 9)) for _ in range(6)]) @@ -592,6 +595,121 @@ class AccountService: return False + @classmethod + def get_admin_through_phone(cls, phone: str): + """ + Get admin account through phone number. + + Returns None if no admin account with this phone number exists. + Raises AccountRegisterError if account is in freeze status. + """ + # We'll check if account is banned or frozen using the phone number + # This implementation assumes phone numbers are unique like emails + + # Find account with end_admin role by phone number + admin_account = ( + db.session.query(Account) + .join(TenantAccountJoin, Account.id == TenantAccountJoin.account_id) + .filter(Account.phone == phone) + .filter(TenantAccountJoin.role == TenantAccountJoinRole.END_ADMIN.value) + .first() + ) + + if not admin_account: + return None + + if admin_account.status == AccountStatus.BANNED.value: + raise Unauthorized("Account is banned.") + + return admin_account + + @classmethod + def is_phone_send_ip_limit(cls, ip_address: str) -> bool: + """ + Check if IP has reached the limit for sending phone verification codes. + Similar to is_email_send_ip_limit but for phone verification. + """ + minute_key = f"phone_send_ip_limit_minute:{ip_address}" + freeze_key = f"phone_send_ip_limit_freeze:{ip_address}" + hour_limit_key = f"phone_send_ip_limit_hour:{ip_address}" + + # check ip is frozen + if redis_client.get(freeze_key): + return True + + # check current minute count + current_minute_count = redis_client.get(minute_key) + if current_minute_count is None: + current_minute_count = 0 + current_minute_count = int(current_minute_count) + + # check current hour count + if current_minute_count > dify_config.EMAIL_SEND_IP_LIMIT_PER_MINUTE: # Use same limit as email + hour_limit_count = redis_client.get(hour_limit_key) + if hour_limit_count is None: + hour_limit_count = 0 + hour_limit_count = int(hour_limit_count) + + if hour_limit_count >= 1: + redis_client.setex(freeze_key, 60 * 60, 1) + return True + else: + redis_client.setex(hour_limit_key, 60 * 10, hour_limit_count + 1) # first time limit 10 minutes + + # add hour limit count + redis_client.incr(hour_limit_key) + redis_client.expire(hour_limit_key, 60 * 60) + + return True + + redis_client.setex(minute_key, 60, current_minute_count + 1) + redis_client.expire(minute_key, 60) + + return False + + @classmethod + def send_phone_code_login(cls, phone: str) -> str: + """ + Send verification code to phone number for admin login. + Returns a token that can be used to verify the code. + """ + if cls.phone_code_login_rate_limiter.is_rate_limited(phone) and not DeploymentConfig().DEBUG: + raise Exception("Phone verification code rate limit exceeded") + + if DeploymentConfig().DEBUG: + # Use a default code for debugging without requiring a config entry + code = dify_config.DEBUG_CODE_FOR_LOGIN + else: + code = "".join([str(random.randint(0, 9)) for _ in range(6)]) + + # Generate a token for code verification + token = TokenManager.generate_token( + token_type="phone_code_login", + account=None, + email=phone, # Using email parameter for phone number + additional_data={"code": code, "phone": phone}, + ) + + # Here you would typically send an SMS with the code + # For now we'll just assume the SMS sending service exists + # send_phone_code_login_sms_task.delay(to=phone, code=code) + + # Log SMS sending in production environment + logging.info(f"Phone verification code sent to {phone}") + + cls.phone_code_login_rate_limiter.increment_rate_limit(phone) + return token + + @classmethod + def get_phone_code_login_data(cls, token: str) -> Optional[dict[str, Any]]: + """Get phone code login data from token""" + return TokenManager.get_token_data(token, "phone_code_login") + + @classmethod + def revoke_phone_code_login_token(cls, token: str) -> None: + """Revoke phone code login token""" + TokenManager.revoke_token(token, "phone_code_login") + def _get_login_cache_key(*, account_id: str, token: str): return f"account_login:{account_id}:{token}" @@ -897,7 +1015,7 @@ class RegisterService: account.last_login_ip = ip_address account.initialized_at = datetime.now(UTC).replace(tzinfo=None) - TenantService.create_owner_tenant_if_not_exist(account=account, is_setup=True) + TenantService.create_owner_tenant_if_not_exist(account=account) dify_setup = DifySetup(version=dify_config.CURRENT_VERSION) db.session.add(dify_setup)