KGPython/controller/RegisterController.py

from robyn import jsonify, Response
from app import app
import os
import uuid
from service.UserService import user_service

# 头像上传目录
AVATAR_UPLOAD_FOLDER = 'resource/avatar'

@app.post("/api/register")
def register_route(request):
    """用户注册接口"""
    try:
        # 获取表单数据
        form_data = getattr(request, 'form_data', {}) if hasattr(request, 'form_data') else {}
        
        # 获取文件上传
        files = getattr(request, 'files', {}) if hasattr(request, 'files') else {}
        
        username = form_data.get("username", "").strip()
        password = form_data.get("password", "").strip()
        
        # 验证必填字段
        if not username or not password:
            return Response(
                status_code=400,
                description=jsonify({"success": False, "message": "用户名和密码不能为空"}),
                headers={"Content-Type": "application/json; charset=utf-8"}
            )
        
        # 检查用户名是否已存在
        existing_user = user_service.get_user_by_username(username)
        if existing_user:
            return Response(
                status_code=409,
                description=jsonify({"success": False, "message": "用户名已存在"}),
                headers={"Content-Type": "application/json; charset=utf-8"}
            )
        
        # 处理头像上传
        avatar_path = "/resource/avatar/4.png"  # 默认头像
        
        # 获取头像文件
        avatar_file = files.get('avatar')
        
        # 如果有头像文件，处理上传
        if avatar_file:
            # 获取文件内容和文件名
            file_content = None
            filename = ""
            
            # Robyn框架中，files字典的值通常是字节数据
            if isinstance(avatar_file, bytes):
                file_content = avatar_file
                filename = "avatar.jpg"
            elif hasattr(avatar_file, 'content'):
                file_content = avatar_file.content
                filename = getattr(avatar_file, 'filename', 'avatar.jpg')
            elif hasattr(avatar_file, 'read'):
                file_content = avatar_file.read()
                filename = getattr(avatar_file, 'filename', 'avatar.jpg')
            else:
                filename = avatar_file if isinstance(avatar_file, str) else 'avatar.jpg'
            
            # 如果我们有文件内容，继续处理
            if file_content:
                # 处理文件名
                file_extension = filename.split('.')[-1] if '.' in filename else 'jpg'
                
                # 验证文件类型
                is_valid_image = False
                
                # 通过扩展名验证
                if file_extension.lower() in ['jpg', 'jpeg', 'png', 'gif', 'bmp', 'webp', 'svg']:
                    is_valid_image = True
                
                # 如果扩展名验证失败，尝试读取文件头验证
                if not is_valid_image:
                    try:
                        if (file_content.startswith(b'\xFF\xD8\xFF') or  # JPEG
                            file_content.startswith(b'\x89PNG\r\n\x1a\n') or  # PNG
                            file_content.startswith(b'GIF87a') or  # GIF
                            file_content.startswith(b'GIF89a')):  # GIF
                            is_valid_image = True
                    except:
                        pass
                
                if not is_valid_image:
                    return Response(
                        status_code=400,
                        description=jsonify({"success": False, "message": "不支持的图片格式"}),
                        headers={"Content-Type": "application/json; charset=utf-8"}
                    )
                
                # 确保上传目录存在
                if not os.path.exists(AVATAR_UPLOAD_FOLDER):
                    os.makedirs(AVATAR_UPLOAD_FOLDER)
                
                # 生成唯一文件名
                unique_filename = f"{uuid.uuid4().hex}.{file_extension}"
                file_path = os.path.join(AVATAR_UPLOAD_FOLDER, unique_filename)
                
                # 保存文件
                try:
                    with open(file_path, 'wb') as f:
                        f.write(file_content)
                    avatar_path = f"/{file_path}"
                except Exception as e:
                    print(f"保存头像失败: {e}")
                    return Response(
                        status_code=500,
                        description=jsonify({"success": False, "message": "头像保存失败"}),
                        headers={"Content-Type": "application/json; charset=utf-8"}
                    )
        
        # 创建用户
        user_id = user_service.create_user(username, password, avatar_path)
        if user_id:
            return Response(
                status_code=201,
                description=jsonify({
                    "success": True, 
                    "message": "注册成功",
                    "user": {
                        "id": user_id,
                        "username": username,
                        "avatar": avatar_path
                    }
                }),
                headers={"Content-Type": "application/json; charset=utf-8"}
            )
        else:
            return Response(
                status_code=500,
                description=jsonify({"success": False, "message": "注册失败，请稍后再试"}),
                headers={"Content-Type": "application/json; charset=utf-8"}
            )
    except Exception as e:
        print(f"注册异常: {e}")
        return Response(
            status_code=500,
            description=jsonify({"success": False, "message": f"注册失败: {str(e)}"}),
            headers={"Content-Type": "application/json; charset=utf-8"}
        )

@app.get("/api/checkUsername")
def check_username_route(request):
    """检查用户名是否可用"""
    try:
        # 获取查询参数
        username = ""
        
        # 尝试从query_params获取
        if hasattr(request, 'query_params') and request.query_params:
            username = request.query_params.get("username", "").strip()
        
        if not username:
            return Response(
                status_code=400,
                description=jsonify({"success": False, "message": "用户名不能为空"}),
                headers={"Content-Type": "application/json; charset=utf-8"}
            )
        
        # 检查用户名是否已存在
        existing_user = user_service.get_user_by_username(username)
        
        return Response(
            status_code=200,
            description=jsonify({
                "success": True, 
                "available": existing_user is None,
                "message": "用户名可用" if not existing_user else "用户名已存在"
            }),
            headers={"Content-Type": "application/json; charset=utf-8"}
        )
    except Exception as e:
        print(f"检查用户名异常: {e}")
        return Response(
            status_code=500,
            description=jsonify({"success": False, "message": f"检查失败: {str(e)}"}),
            headers={"Content-Type": "application/json; charset=utf-8"}
        )