90 lines
2.7 KiB
Python
90 lines
2.7 KiB
Python
from datetime import date
|
|
from models import DailyParseStat, UserGroup
|
|
from models import db
|
|
|
|
class RateLimiter:
|
|
"""限流器"""
|
|
|
|
@staticmethod
|
|
def check_limit(user_id=None, ip_address=None):
|
|
"""检查是否超过限制"""
|
|
today = date.today()
|
|
|
|
if user_id:
|
|
# 已登录用户
|
|
stat = DailyParseStat.query.filter_by(
|
|
user_id=user_id,
|
|
date=today
|
|
).first()
|
|
|
|
from models import User
|
|
user = User.query.get(user_id)
|
|
group = UserGroup.query.get(user.group_id)
|
|
limit = group.daily_limit
|
|
|
|
current_count = stat.parse_count if stat else 0
|
|
|
|
return {
|
|
'allowed': current_count < limit,
|
|
'current': current_count,
|
|
'limit': limit,
|
|
'remaining': max(0, limit - current_count)
|
|
}
|
|
else:
|
|
# 游客
|
|
stat = DailyParseStat.query.filter_by(
|
|
ip_address=ip_address,
|
|
date=today
|
|
).first()
|
|
|
|
from models import SiteConfig
|
|
config = SiteConfig.query.filter_by(config_key='guest_daily_limit').first()
|
|
limit = int(config.config_value) if config else 5
|
|
|
|
current_count = stat.parse_count if stat else 0
|
|
|
|
return {
|
|
'allowed': current_count < limit,
|
|
'current': current_count,
|
|
'limit': limit,
|
|
'remaining': max(0, limit - current_count)
|
|
}
|
|
|
|
@staticmethod
|
|
def increment_count(user_id=None, ip_address=None, success=True):
|
|
"""增加计数"""
|
|
today = date.today()
|
|
|
|
if user_id:
|
|
stat = DailyParseStat.query.filter_by(
|
|
user_id=user_id,
|
|
date=today
|
|
).first()
|
|
|
|
if not stat:
|
|
stat = DailyParseStat(user_id=user_id, date=today)
|
|
db.session.add(stat)
|
|
|
|
stat.parse_count = (stat.parse_count or 0) + 1
|
|
if success:
|
|
stat.success_count = (stat.success_count or 0) + 1
|
|
else:
|
|
stat.fail_count = (stat.fail_count or 0) + 1
|
|
else:
|
|
stat = DailyParseStat.query.filter_by(
|
|
ip_address=ip_address,
|
|
date=today
|
|
).first()
|
|
|
|
if not stat:
|
|
stat = DailyParseStat(ip_address=ip_address, date=today)
|
|
db.session.add(stat)
|
|
|
|
stat.parse_count = (stat.parse_count or 0) + 1
|
|
if success:
|
|
stat.success_count = (stat.success_count or 0) + 1
|
|
else:
|
|
stat.fail_count = (stat.fail_count or 0) + 1
|
|
|
|
db.session.commit()
|