122 lines
3.8 KiB
Python
122 lines
3.8 KiB
Python
import sqlite3
|
|
import os
|
|
from datetime import datetime, timedelta
|
|
from config import DATABASE, DAILY_TRAFFIC_LIMIT
|
|
|
|
def get_db():
|
|
conn = sqlite3.connect(DATABASE)
|
|
conn.row_factory = sqlite3.Row
|
|
return conn
|
|
|
|
def init_db():
|
|
os.makedirs(os.path.dirname(DATABASE) if os.path.dirname(DATABASE) else '.', exist_ok=True)
|
|
conn = get_db()
|
|
conn.execute('''
|
|
CREATE TABLE IF NOT EXISTS files (
|
|
id TEXT PRIMARY KEY,
|
|
filename TEXT NOT NULL,
|
|
filepath TEXT NOT NULL,
|
|
filesize INTEGER NOT NULL,
|
|
expiry_hours INTEGER NOT NULL,
|
|
created_at TIMESTAMP NOT NULL,
|
|
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()
|
|
|
|
def add_file(file_id, filename, filepath, filesize, expiry_hours):
|
|
now = datetime.utcnow()
|
|
expires = now + timedelta(hours=expiry_hours)
|
|
conn = get_db()
|
|
conn.execute(
|
|
'INSERT INTO files (id, filename, filepath, filesize, expiry_hours, created_at, expires_at) VALUES (?, ?, ?, ?, ?, ?, ?)',
|
|
(file_id, filename, filepath, filesize, expiry_hours, now, expires)
|
|
)
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def get_file(file_id):
|
|
conn = get_db()
|
|
row = conn.execute('SELECT * FROM files WHERE id = ?', (file_id,)).fetchone()
|
|
conn.close()
|
|
return row
|
|
|
|
def delete_file(file_id):
|
|
conn = get_db()
|
|
conn.execute('DELETE FROM files WHERE id = ?', (file_id,))
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def cleanup_expired():
|
|
now = datetime.utcnow()
|
|
conn = get_db()
|
|
expired = conn.execute('SELECT * FROM files WHERE expires_at < ?', (now,)).fetchall()
|
|
for row in expired:
|
|
if os.path.exists(row['filepath']):
|
|
os.remove(row['filepath'])
|
|
conn.execute('DELETE FROM files WHERE expires_at < ?', (now,))
|
|
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
|