From f5727ec176f58d04add5da87d558e9e206d784b3 Mon Sep 17 00:00:00 2001 From: OpenCode Bot Date: Fri, 1 May 2026 09:02:39 +0800 Subject: [PATCH] feat(traffic): add per-IP daily traffic tracking with 20GB limit for uploads and downloads --- app.py | 30 +++++++++++++++++++++++--- config.py | 2 ++ database.py | 62 ++++++++++++++++++++++++++++++++++++++++++++++++++++- 3 files changed, 90 insertions(+), 4 deletions(-) diff --git a/app.py b/app.py index 7a75657..fb51432 100644 --- a/app.py +++ b/app.py @@ -2,8 +2,8 @@ import os import uuid from datetime import datetime from flask import Flask, request, render_template, send_file, jsonify, url_for, abort -from config import UPLOAD_FOLDER, SECRET_KEY, MAX_CONTENT_LENGTH, EXPIRY_OPTIONS -from database import init_db, add_file, get_file, delete_file, cleanup_expired +from config import UPLOAD_FOLDER, SECRET_KEY, MAX_CONTENT_LENGTH, EXPIRY_OPTIONS, DAILY_TRAFFIC_LIMIT +from database import init_db, add_file, get_file, delete_file, cleanup_expired, add_upload_traffic, add_download_traffic, get_client_ip, is_traffic_exceeded, get_daily_traffic app = Flask(__name__) app.config['SECRET_KEY'] = SECRET_KEY @@ -30,6 +30,12 @@ def upload(): expiry_seconds = EXPIRY_OPTIONS.get(expiry_key, EXPIRY_OPTIONS['24h']) expiry_hours = expiry_seconds // 3600 + ip = get_client_ip(request) + content_length = request.content_length or 0 + + if is_traffic_exceeded(ip, content_length, 'upload'): + return jsonify({'error': 'Daily traffic limit exceeded (20GB)'}), 429 + file_id = str(uuid.uuid4()) filename = file.filename filepath = os.path.join(UPLOAD_FOLDER, file_id) @@ -37,6 +43,7 @@ def upload(): filesize = os.path.getsize(filepath) add_file(file_id, filename, filepath, filesize, expiry_hours) + add_upload_traffic(ip, filesize) share_url = url_for('download_file', file_id=file_id, _external=True) return jsonify({'id': file_id, 'filename': filename, 'share_url': share_url}) @@ -54,6 +61,12 @@ def api_upload(): expiry_seconds = EXPIRY_OPTIONS.get(expiry_key, EXPIRY_OPTIONS['24h']) expiry_hours = expiry_seconds // 3600 + ip = get_client_ip(request) + content_length = request.content_length or 0 + + if is_traffic_exceeded(ip, content_length, 'upload'): + return jsonify({'error': 'Daily traffic limit exceeded (20GB)'}), 429 + file_id = str(uuid.uuid4()) filename = file.filename filepath = os.path.join(UPLOAD_FOLDER, file_id) @@ -61,6 +74,7 @@ def api_upload(): filesize = os.path.getsize(filepath) add_file(file_id, filename, filepath, filesize, expiry_hours) + add_upload_traffic(ip, filesize) share_url = url_for('download_file', file_id=file_id, _external=True) return jsonify({'id': file_id, 'filename': filename, 'filesize': filesize, 'expiry_hours': expiry_hours, 'share_url': share_url}) @@ -79,12 +93,18 @@ def api_get_file(file_id): row = get_file(file_id) if not row: return jsonify({'error': 'File not found or expired'}), 404 + + ip = get_client_ip(request) + upload, download = get_daily_traffic(ip) return jsonify({ 'id': row['id'], 'filename': row['filename'], 'filesize': row['filesize'], 'created_at': row['created_at'], - 'expires_at': row['expires_at'] + 'expires_at': row['expires_at'], + 'daily_upload': upload, + 'daily_download': download, + 'traffic_limit': DAILY_TRAFFIC_LIMIT }) @app.route('/download/') @@ -92,6 +112,10 @@ def serve_file(file_id): row = get_file(file_id) if not row: abort(404) + + ip = get_client_ip(request) + add_download_traffic(ip, row['filesize']) + return send_file(row['filepath'], download_name=row['filename'], as_attachment=True) if __name__ == '__main__': diff --git a/config.py b/config.py index cc7efec..13739e7 100644 --- a/config.py +++ b/config.py @@ -7,6 +7,8 @@ UPLOAD_FOLDER = os.path.join(BASE_DIR, 'uploads') DATABASE = os.path.join(BASE_DIR, 'files.db') MAX_CONTENT_LENGTH = 500 * 1024 * 1024 +DAILY_TRAFFIC_LIMIT = 20 * 1024 * 1024 * 1024 + EXPIRY_OPTIONS = { '1h': 60 * 60, '24h': 24 * 60 * 60, diff --git a/database.py b/database.py index 8372439..4b65dbf 100644 --- a/database.py +++ b/database.py @@ -1,7 +1,7 @@ import sqlite3 import os from datetime import datetime, timedelta -from config import DATABASE +from config import DATABASE, DAILY_TRAFFIC_LIMIT def get_db(): conn = sqlite3.connect(DATABASE) @@ -22,6 +22,16 @@ def init_db(): expires_at TIMESTAMP NOT NULL ) ''') + conn.execute(''' + CREATE TABLE IF NOT EXISTS ip_traffic ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + ip TEXT NOT NULL, + date TEXT NOT NULL, + upload_bytes INTEGER NOT NULL DEFAULT 0, + download_bytes INTEGER NOT NULL DEFAULT 0, + UNIQUE(ip, date) + ) + ''') conn.commit() conn.close() @@ -59,3 +69,53 @@ def cleanup_expired(): conn.commit() conn.close() return len(expired) + +def get_client_ip(request): + if 'X-Forwarded-For' in request.headers: + return request.headers['X-Forwarded-For'].split(',')[0].strip() + if 'X-Real-IP' in request.headers: + return request.headers['X-Real-IP'].strip() + return request.remote_addr + +def add_upload_traffic(ip, bytes_count): + today = datetime.utcnow().strftime('%Y-%m-%d') + conn = get_db() + conn.execute(''' + INSERT INTO ip_traffic (ip, date, upload_bytes, download_bytes) + VALUES (?, ?, ?, 0) + ON CONFLICT(ip, date) DO UPDATE SET upload_bytes = upload_bytes + ? + ''', (ip, today, bytes_count, bytes_count)) + conn.commit() + conn.close() + +def add_download_traffic(ip, bytes_count): + today = datetime.utcnow().strftime('%Y-%m-%d') + conn = get_db() + conn.execute(''' + INSERT INTO ip_traffic (ip, date, upload_bytes, download_bytes) + VALUES (?, ?, 0, ?) + ON CONFLICT(ip, date) DO UPDATE SET download_bytes = download_bytes + ? + ''', (ip, today, bytes_count, bytes_count)) + conn.commit() + conn.close() + +def get_daily_traffic(ip): + today = datetime.utcnow().strftime('%Y-%m-%d') + conn = get_db() + row = conn.execute( + 'SELECT upload_bytes, download_bytes FROM ip_traffic WHERE ip = ? AND date = ?', + (ip, today) + ).fetchone() + conn.close() + if row: + return row['upload_bytes'], row['download_bytes'] + return 0, 0 + +def is_traffic_exceeded(ip, additional_bytes, direction='upload'): + upload, download = get_daily_traffic(ip) + total = upload + download + if direction == 'upload': + total += additional_bytes + else: + total += additional_bytes + return total > DAILY_TRAFFIC_LIMIT