DataOps-platform/app/core/system/auth.py

"""
系统用户认证模块
提供用户注册、登录验证等功能
"""

import base64
import logging
import time
import uuid
from functools import wraps
from urllib.parse import unquote, urlparse

import psycopg2
from flask import current_app, jsonify, request

logger = logging.getLogger(__name__)

# PostgreSQL连接池
pg_pool = None


def get_pg_connection():
    """
    获取PostgreSQL数据库连接

    Returns:
        connection: PostgreSQL连接对象
    """
    global pg_pool

    if pg_pool is None:
        try:
            # 解析SQLAlchemy URI，处理包含特殊字符的密码
            db_uri = current_app.config["SQLALCHEMY_DATABASE_URI"]

            # 尝试使用urlparse解析
            uri = urlparse(db_uri)

            # 如果解析失败（缺少用户名或主机名）或密码包含特殊字符导致解析错误，使用手动解析
            if uri.username is None or uri.hostname is None:
                # 手动解析URI: postgresql://username:password@host:port/database
                scheme_end = db_uri.find("://")
                if scheme_end == -1:
                    raise ValueError("Invalid database URI format")

                auth_and_host = db_uri[scheme_end + 3 :]  # 跳过 '://'
                at_pos = auth_and_host.rfind("@")  # 从右向左查找最后一个@

                if at_pos == -1:
                    raise ValueError("Invalid database URI: missing @ separator")

                auth_part = auth_and_host[:at_pos]
                host_part = auth_and_host[at_pos + 1 :]

                # 解析用户名和密码（可能包含特殊字符）
                colon_pos = auth_part.find(":")
                if colon_pos == -1:
                    username = unquote(auth_part)
                    password = None
                else:
                    username = unquote(auth_part[:colon_pos])
                    password = unquote(auth_part[colon_pos + 1 :])

                # 解析主机、端口和数据库
                slash_pos = host_part.find("/")
                if slash_pos == -1:
                    raise ValueError("Invalid database URI: missing database name")

                host_port = host_part[:slash_pos]
                database = unquote(host_part[slash_pos + 1 :])

                # 解析主机和端口
                colon_pos = host_port.find(":")
                if colon_pos == -1:
                    hostname = host_port
                    port = 5432
                else:
                    hostname = host_port[:colon_pos]
                    port = int(host_port[colon_pos + 1 :])
            else:
                # urlparse解析成功，解码可能被URL编码的字段
                username = unquote(uri.username) if uri.username else None
                password = unquote(uri.password) if uri.password else None
                database = (
                    unquote(uri.path[1:]) if uri.path and len(uri.path) > 1 else None
                )
                hostname = uri.hostname
                port = uri.port or 5432

            # 验证必需的字段（username, database, hostname 是必需的，password 是可选的）
            if not all([username, database, hostname]):
                raise ValueError(
                    "Missing required database connection parameters: username, database, and hostname are required"
                )

            # 创建连接池
            pg_pool = psycopg2.pool.SimpleConnectionPool(  # type: ignore[attr-defined]
                1,
                20,
                host=hostname,
                database=database,
                user=username,
                password=password,
                port=str(port),
            )
            logger.info("PostgreSQL连接池初始化成功")
        except Exception as e:
            logger.error(f"PostgreSQL连接池初始化失败: {str(e)}")
            raise

    return pg_pool.getconn()


def release_pg_connection(conn):
    """
    释放PostgreSQL连接到连接池

    Args:
        conn: 数据库连接对象
    """
    global pg_pool
    if pg_pool and conn:
        pg_pool.putconn(conn)


def encode_password(password):
    """
    对密码进行base64编码

    Args:
        password: 原始密码

    Returns:
        str: 编码后的密码
    """
    return base64.b64encode(password.encode("utf-8")).decode("utf-8")


def create_user_table():
    """
    创建用户表，如果不存在

    Returns:
        bool: 是否成功创建
    """
    conn = None
    try:
        conn = get_pg_connection()
        cursor = conn.cursor()

        # 创建用户表
        create_table_query = """
        CREATE TABLE IF NOT EXISTS users (
            id VARCHAR(100) PRIMARY KEY,
            username VARCHAR(50) UNIQUE NOT NULL,
            password VARCHAR(100) NOT NULL,
            created_at FLOAT NOT NULL,
            last_login FLOAT,
            is_admin BOOLEAN DEFAULT FALSE
        );
        """
        cursor.execute(create_table_query)

        # 创建索引加速查询
        create_index_query = """
        CREATE INDEX IF NOT EXISTS idx_users_username ON users(username);
        """
        cursor.execute(create_index_query)

        conn.commit()
        cursor.close()

        logger.info("用户表创建成功")
        return True
    except Exception as e:
        logger.error(f"创建用户表失败: {str(e)}")
        if conn:
            conn.rollback()
        return False
    finally:
        if conn:
            release_pg_connection(conn)


def register_user(username, password):
    """
    注册新用户

    Args:
        username: 用户名
        password: 密码

    Returns:
        tuple: (是否成功, 消息)
    """
    conn = None
    try:
        # 确保表已创建
        create_user_table()

        # 对密码进行编码
        encoded_password = encode_password(password)

        # 生成用户ID
        user_id = str(uuid.uuid4())

        conn = get_pg_connection()
        cursor = conn.cursor()

        # 检查用户名是否存在
        check_query = "SELECT username FROM users WHERE username = %s"
        cursor.execute(check_query, (username,))

        if cursor.fetchone():
            return False, "用户名已存在"

        # 创建用户
        insert_query = """
        INSERT INTO users (id, username, password, created_at, last_login)
        VALUES (%s, %s, %s, %s, %s)
        """
        cursor.execute(
            insert_query, (user_id, username, encoded_password, time.time(), None)
        )

        conn.commit()
        cursor.close()

        return True, "注册成功"
    except Exception as e:
        logger.error(f"用户注册失败: {str(e)}")
        if conn:
            conn.rollback()
        return False, f"注册失败: {str(e)}"
    finally:
        if conn:
            release_pg_connection(conn)


def login_user(username, password):
    """
    用户登录验证

    Args:
        username: 用户名
        password: 密码

    Returns:
        tuple: (是否成功, 用户信息/错误消息)
    """
    conn = None
    try:
        # 对输入的密码进行编码
        encoded_password = encode_password(password)

        conn = get_pg_connection()
        cursor = conn.cursor()

        # 查询用户
        query = """
        SELECT id, username, password, created_at, last_login, is_admin
        FROM users WHERE username = %s
        """
        cursor.execute(query, (username,))

        user = cursor.fetchone()

        # 检查用户是否存在
        if not user:
            return False, "用户名或密码错误"

        # 验证密码
        if user[2] != encoded_password:
            return False, "用户名或密码错误"

        # 更新最后登录时间
        current_time = time.time()
        update_query = """
        UPDATE users SET last_login = %s WHERE username = %s
        """
        cursor.execute(update_query, (current_time, username))

        conn.commit()

        # 构建用户信息
        user_info = {
            "id": user[0],
            "username": user[1],
            "created_at": user[3],
            "last_login": current_time,
            "is_admin": user[5] if len(user) > 5 else False,
        }

        cursor.close()

        return True, user_info
    except Exception as e:
        logger.error(f"用户登录失败: {str(e)}")
        if conn:
            conn.rollback()
        return False, f"登录失败: {str(e)}"
    finally:
        if conn:
            release_pg_connection(conn)


def get_user_by_username(username):
    """
    根据用户名获取用户信息

    Args:
        username: 用户名

    Returns:
        dict: 用户信息(不包含密码)
    """
    conn = None
    try:
        conn = get_pg_connection()
        cursor = conn.cursor()

        query = """
        SELECT id, username, created_at, last_login, is_admin
        FROM users WHERE username = %s
        """
        cursor.execute(query, (username,))

        user = cursor.fetchone()
        cursor.close()

        if not user:
            return None

        user_info = {
            "id": user[0],
            "username": user[1],
            "created_at": user[2],
            "last_login": user[3],
            "is_admin": user[4] if user[4] is not None else False,
        }

        return user_info
    except Exception as e:
        logger.error(f"获取用户信息失败: {str(e)}")
        return None
    finally:
        if conn:
            release_pg_connection(conn)


def init_db():
    """
    初始化数据库，创建用户表

    Returns:
        bool: 是否成功初始化
    """
    return create_user_table()


def require_auth(f):
    @wraps(f)
    def decorated(*args, **kwargs):
        auth_header = request.headers.get("Authorization")
        if not auth_header:
            return jsonify({"message": "缺少认证头"}), 401

        try:
            # 验证认证头
            if auth_header != current_app.config["SECRET_KEY"]:
                return jsonify({"message": "无效的认证信息"}), 401

            return f(*args, **kwargs)
        except Exception:
            return jsonify({"message": "认证失败"}), 401

    return decorated