temp-file-trans/database.py

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

def get_monthly_stats():
    now = datetime.utcnow()
    month_start = now.strftime('%Y-%m-01')
    conn = get_db()
    traffic_row = conn.execute(
        'SELECT COALESCE(SUM(upload_bytes),0) AS up, COALESCE(SUM(download_bytes),0) AS down FROM ip_traffic WHERE date >= ?',
        (month_start,)
    ).fetchone()
    ip_count = conn.execute(
        'SELECT COUNT(DISTINCT ip) FROM ip_traffic WHERE date >= ?',
        (month_start,)
    ).fetchone()[0]
    file_count = conn.execute(
        'SELECT COUNT(*) FROM files WHERE created_at >= ?',
        (month_start,)
    ).fetchone()[0]
    conn.close()
    total = (traffic_row['up'] or 0) + (traffic_row['down'] or 0)
    return {
        'total_bytes': total,
        'total_gb': round(total / (1024**3), 2),
        'ip_count': ip_count,
        'file_count': file_count,
    }