RamEx / api /utils /auth.py
zdy10046's picture
deploy RamEx to Hugging Face without binary files
e657e99
Raw
History Blame Contribute Delete
9.91 kB
import requests
from rest_framework.authentication import BaseAuthentication
from rest_framework.exceptions import AuthenticationFailed
from django.contrib.auth.models import User
from django.conf import settings
import logging
import os
from pathlib import Path
logger = logging.getLogger(__name__)
class ThirdPartyUserManager:
"""第三方用户管理工具类"""
@classmethod
def get_or_create_user_from_third_party(cls, third_party_username):
"""
根据第三方用户名获取或创建本地用户
"""
from api.models import ThirdPartyUserMapping
# 首先尝试通过用户名查找现有用户映射
try:
mapping = ThirdPartyUserMapping.objects.get(third_party_username=third_party_username)
return mapping.user
except ThirdPartyUserMapping.DoesNotExist:
pass
# 如果没找到映射,创建新的本地用户和映射
# 使用第三方用户名作为username,确保唯一性
username = f"tp_{third_party_username}"
# 确保用户名唯一
if User.objects.filter(username=username).exists():
# 如果存在同名用户,添加后缀
counter = 1
base_username = username
while User.objects.filter(username=username).exists():
username = f"{base_username}_{counter}"
counter += 1
# 创建本地用户
user = User.objects.create(
username=username,
email=f"{third_party_username}@thirdparty.com",
first_name=third_party_username,
)
# 创建映射关系(使用用户名作为唯一标识)
ThirdPartyUserMapping.objects.create(
user=user,
third_party_user_id=hash(third_party_username) % 2147483647, # 生成一个伪ID
third_party_username=third_party_username
)
# 创建用户数据目录结构
cls.create_user_directories(user)
logger.info(f"Created new local user {username} for third party user {third_party_username}")
return user
@classmethod
def get_user_by_third_party_username(cls, third_party_username):
"""
根据第三方用户名获取本地用户
"""
from api.models import ThirdPartyUserMapping
try:
mapping = ThirdPartyUserMapping.objects.get(third_party_username=third_party_username)
return mapping.user
except ThirdPartyUserMapping.DoesNotExist:
return None
@classmethod
def create_user_directories(cls, user):
"""
为用户创建必要的数据目录结构
"""
try:
# 基础媒体目录
media_root = Path(settings.MEDIA_ROOT)
# 用户目录结构
user_base_dir = media_root / 'users' / str(user.id)
projects_dir = user_base_dir / 'projects'
# 创建目录
user_base_dir.mkdir(parents=True, exist_ok=True)
projects_dir.mkdir(parents=True, exist_ok=True)
logger.info(f"Created directory structure for user {user.username} at {user_base_dir}")
except Exception as e:
logger.error(f"Failed to create directories for user {user.username}: {e}")
@classmethod
def get_user_by_third_party_id(cls, third_party_user_id):
"""
根据第三方用户ID获取本地用户
"""
from api.models import ThirdPartyUserMapping
try:
mapping = ThirdPartyUserMapping.objects.get(third_party_user_id=third_party_user_id)
return mapping.user
except ThirdPartyUserMapping.DoesNotExist:
return None
class ThirdPartyAuth:
"""第三方认证工具类"""
BASE_URL = "http://27.223.62.254:10008"
@classmethod
def login(cls, username, password):
"""
登录获取token
"""
try:
url = f"{cls.BASE_URL}/login/{username},{password},1"
headers = {
'accept': 'application/json',
'X-CSRFToken': 'ARNjGjIoXRiA5I5LwmKcUDB5GGzPtMofGDJW5UcBsjWAolmvA8iZoLjWMxFepakY'
}
response = requests.get(url, headers=headers, timeout=10)
if response.status_code == 200:
data = response.json()
if 'token' in data:
return {
'success': True,
'token': data['token'],
'user_id': data.get('user_id'),
'username': data.get('username'),
'message': data.get('msg', '')
}
return {
'success': False,
'message': f'登录失败: HTTP {response.status_code}'
}
except requests.RequestException as e:
logger.error(f"Third party login error: {e}")
return {
'success': False,
'message': f'登录服务连接失败: {str(e)}'
}
@classmethod
def validate_token(cls, token):
"""
验证token有效性
"""
try:
url = f"{cls.BASE_URL}/user/valid_token/"
headers = {
'accept': 'application/json',
'X-CSRFToken': 'ARNjGjIoXRiA5I5LwmKcUDB5GGzPtMofGDJW5UcBsjWAolmvA8iZoLjWMxFepakY'
}
files = {'user_token': (None, token)}
response = requests.post(url, headers=headers, files=files, timeout=10)
if response.status_code == 200:
data = response.json()
if data.get('code') == '200' and data.get('message') == 'succeed':
return {
'valid': True,
'data': data.get('data', {}),
'user_name': data.get('data', {}).get('user_name')
}
return {
'valid': False,
'message': f'Token验证失败: HTTP {response.status_code}'
}
except requests.RequestException as e:
logger.error(f"Token validation error: {e}")
return {
'valid': False,
'message': f'Token验证服务连接失败: {str(e)}'
}
@classmethod
def get_user_info_by_token(cls, token):
"""
通过token获取用户信息
"""
validation_result = cls.validate_token(token)
if validation_result['valid']:
return validation_result['data']
return None
class ThirdPartyTokenAuthentication(BaseAuthentication):
"""
第三方Token认证类
只接受Authorization header中的token进行认证
"""
def authenticate(self, request):
"""
验证请求中的token
"""
# 从请求头获取token
token = self.get_token_from_request(request)
if not token:
return None
# 验证token并获取用户
user = self.get_or_create_user(token)
if not user:
raise AuthenticationFailed('Invalid token')
return (user, token)
def get_token_from_request(self, request):
"""
从请求中获取token
只支持Authorization header方式:
- Authorization: Token <token>
- Authorization: Bearer <token>
"""
# 从Authorization header获取
auth_header = request.META.get('HTTP_AUTHORIZATION')
if not auth_header:
return None
# 支持Token格式
if auth_header.startswith('Token '):
return auth_header[6:]
# 支持Bearer格式
if auth_header.startswith('Bearer '):
return auth_header[7:]
# 不支持其他格式
return None
def get_or_create_user(self, token):
"""
获取或创建用户
"""
from api.models import TokenStorage
# 首先尝试从TokenStorage中获取用户信息
try:
token_storage = TokenStorage.objects.get(token=token)
# 更新最后使用时间
token_storage.save() # 这会触发last_used字段更新
return token_storage.user
except TokenStorage.DoesNotExist:
pass
# 如果TokenStorage中没有记录,验证token有效性
validation_result = ThirdPartyAuth.validate_token(token)
if not validation_result['valid']:
logger.warning(f"Invalid token: {token[:8]}...")
return None
# 获取用户名
user_name = validation_result.get('user_name')
if not user_name:
logger.error(f"No user_name in validation result for token: {token[:8]}...")
return None
# 根据用户名获取或创建用户
user = ThirdPartyUserManager.get_or_create_user_from_third_party(user_name)
# 将token存储到TokenStorage中
try:
TokenStorage.objects.update_or_create(
token=token,
defaults={
'user': user,
'third_party_user_id': hash(user_name) % 2147483647,
'third_party_username': user_name,
}
)
logger.info(f"Token stored for user: {user_name}")
except Exception as e:
logger.error(f"Failed to store token: {e}")
return user