init

utils/limiter.py

@@ -0,0 +1,89 @@
+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()