mirror of
https://github.com/mattintech/simplefileupload-server.git
synced 2026-07-11 17:41:53 +00:00
Migrate storage from JSON to SQLite with enhanced MVC architecture
- Replace JSON file storage with SQLite database for improved concurrency and data integrity - Implement repository pattern with dedicated model layer (admin_model, file_model, chunk_model) - Add database.py with automatic migration from existing JSON files - Enable WAL mode for thread-safe concurrent access across Gunicorn workers - Store database files in dedicated db/ folder - Update controllers to use model layer instead of direct JSON access - Add admin authentication system with TOTP 2FA support
This commit is contained in:
@@ -0,0 +1,35 @@
|
||||
from .file_metadata import FileMetadata
|
||||
|
||||
# Database initialization
|
||||
from .database import init_db, get_db_connection, migrate_from_json
|
||||
|
||||
# Admin model functions
|
||||
from .admin_model import (
|
||||
get_admin,
|
||||
create_admin,
|
||||
admin_exists as model_admin_exists,
|
||||
update_last_login
|
||||
)
|
||||
|
||||
# File model functions
|
||||
from .file_model import (
|
||||
create_file,
|
||||
get_file,
|
||||
increment_download,
|
||||
delete_file as model_delete_file,
|
||||
get_expired_files,
|
||||
cleanup_expired,
|
||||
get_all_files
|
||||
)
|
||||
|
||||
# Chunk model functions
|
||||
from .chunk_model import (
|
||||
create_session,
|
||||
get_session,
|
||||
add_chunk,
|
||||
get_received_chunks,
|
||||
delete_session,
|
||||
get_abandoned_sessions,
|
||||
cleanup_abandoned,
|
||||
update_received_chunks
|
||||
)
|
||||
|
||||
91
src/models/admin_model.py
Normal file
91
src/models/admin_model.py
Normal file
@@ -0,0 +1,91 @@
|
||||
import sqlite3
|
||||
from models.database import get_db
|
||||
import utils.mLogger as log
|
||||
|
||||
def get_admin():
|
||||
"""Get admin user data
|
||||
|
||||
Returns:
|
||||
dict: Admin data with keys: username, password_hash, totp_secret, created_at, last_login
|
||||
None: If no admin exists
|
||||
"""
|
||||
try:
|
||||
with get_db() as conn:
|
||||
cursor = conn.cursor()
|
||||
row = cursor.execute("""
|
||||
SELECT username, password_hash, totp_secret, created_at, last_login
|
||||
FROM admin
|
||||
LIMIT 1
|
||||
""").fetchone()
|
||||
|
||||
if row:
|
||||
return dict(row)
|
||||
return None
|
||||
except sqlite3.Error as e:
|
||||
log.e(f"Error getting admin data: {e}")
|
||||
return None
|
||||
|
||||
def create_admin(username, password_hash, totp_secret):
|
||||
"""Create admin user
|
||||
|
||||
Args:
|
||||
username: Admin username
|
||||
password_hash: Hashed password
|
||||
totp_secret: TOTP secret for 2FA
|
||||
|
||||
Returns:
|
||||
bool: True if successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
with get_db() as conn:
|
||||
from datetime import datetime
|
||||
|
||||
conn.execute("""
|
||||
INSERT INTO admin (username, password_hash, totp_secret, created_at)
|
||||
VALUES (?, ?, ?, ?)
|
||||
""", (username, password_hash, totp_secret, datetime.now().isoformat()))
|
||||
|
||||
log.i(f"Admin user created: {username}")
|
||||
return True
|
||||
except sqlite3.Error as e:
|
||||
log.e(f"Error creating admin: {e}")
|
||||
return False
|
||||
|
||||
def admin_exists():
|
||||
"""Check if an admin user has been created
|
||||
|
||||
Returns:
|
||||
bool: True if admin exists, False otherwise
|
||||
"""
|
||||
try:
|
||||
with get_db() as conn:
|
||||
cursor = conn.cursor()
|
||||
count = cursor.execute("SELECT COUNT(*) FROM admin").fetchone()[0]
|
||||
return count > 0
|
||||
except sqlite3.Error as e:
|
||||
log.e(f"Error checking if admin exists: {e}")
|
||||
return False
|
||||
|
||||
def update_last_login(username):
|
||||
"""Update admin's last login time
|
||||
|
||||
Args:
|
||||
username: Admin username
|
||||
|
||||
Returns:
|
||||
bool: True if successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
with get_db() as conn:
|
||||
from datetime import datetime
|
||||
|
||||
conn.execute("""
|
||||
UPDATE admin
|
||||
SET last_login = ?
|
||||
WHERE username = ?
|
||||
""", (datetime.now().isoformat(), username))
|
||||
|
||||
return True
|
||||
except sqlite3.Error as e:
|
||||
log.e(f"Error updating last login: {e}")
|
||||
return False
|
||||
262
src/models/chunk_model.py
Normal file
262
src/models/chunk_model.py
Normal file
@@ -0,0 +1,262 @@
|
||||
import sqlite3
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
from datetime import datetime, timedelta
|
||||
from models.database import get_db
|
||||
import utils.mLogger as log
|
||||
|
||||
def create_session(upload_id, filename, total_chunks, chunk_dir, max_downloads, expiry_hours):
|
||||
"""Create a new chunked upload session
|
||||
|
||||
Args:
|
||||
upload_id: 16-character unique upload ID
|
||||
filename: Original filename
|
||||
total_chunks: Total number of chunks expected
|
||||
chunk_dir: Directory where chunks are stored
|
||||
max_downloads: Maximum downloads for final file
|
||||
expiry_hours: Expiry hours for final file
|
||||
|
||||
Returns:
|
||||
bool: True if successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
with get_db() as conn:
|
||||
conn.execute("""
|
||||
INSERT INTO chunk_sessions
|
||||
(upload_id, filename, total_chunks, received_chunks, created_at,
|
||||
chunk_dir, max_downloads, expiry_hours)
|
||||
VALUES (?, ?, ?, '[]', ?, ?, ?, ?)
|
||||
""", (upload_id, filename, total_chunks, datetime.now().isoformat(),
|
||||
chunk_dir, max_downloads, expiry_hours))
|
||||
|
||||
log.i(f"Created chunk session: {upload_id}")
|
||||
return True
|
||||
except sqlite3.Error as e:
|
||||
log.e(f"Error creating chunk session: {e}")
|
||||
return False
|
||||
|
||||
def get_session(upload_id):
|
||||
"""Get chunk upload session data
|
||||
|
||||
Args:
|
||||
upload_id: 16-character upload ID
|
||||
|
||||
Returns:
|
||||
dict: Session data with keys: upload_id, filename, total_chunks, received_chunks (list),
|
||||
created_at, chunk_dir, max_downloads, expiry_hours
|
||||
None: If session not found
|
||||
"""
|
||||
try:
|
||||
with get_db() as conn:
|
||||
cursor = conn.cursor()
|
||||
row = cursor.execute("""
|
||||
SELECT upload_id, filename, total_chunks, received_chunks, created_at,
|
||||
chunk_dir, max_downloads, expiry_hours
|
||||
FROM chunk_sessions
|
||||
WHERE upload_id = ?
|
||||
""", (upload_id,)).fetchone()
|
||||
|
||||
if row:
|
||||
data = dict(row)
|
||||
# Parse received_chunks from JSON string to list
|
||||
data['received_chunks'] = json.loads(data['received_chunks'])
|
||||
return data
|
||||
return None
|
||||
except sqlite3.Error as e:
|
||||
log.e(f"Error getting chunk session {upload_id}: {e}")
|
||||
return None
|
||||
|
||||
def add_chunk(upload_id, chunk_index):
|
||||
"""Add a chunk index to the received chunks list
|
||||
|
||||
Args:
|
||||
upload_id: 16-character upload ID
|
||||
chunk_index: Index of the chunk received
|
||||
|
||||
Returns:
|
||||
bool: True if successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
with get_db() as conn:
|
||||
cursor = conn.cursor()
|
||||
|
||||
# Get current received_chunks
|
||||
row = cursor.execute("""
|
||||
SELECT received_chunks FROM chunk_sessions WHERE upload_id = ?
|
||||
""", (upload_id,)).fetchone()
|
||||
|
||||
if not row:
|
||||
log.e(f"Chunk session not found: {upload_id}")
|
||||
return False
|
||||
|
||||
received = json.loads(row['received_chunks'])
|
||||
|
||||
# Add chunk index if not already present
|
||||
if chunk_index not in received:
|
||||
received.append(chunk_index)
|
||||
|
||||
# Update database
|
||||
cursor.execute("""
|
||||
UPDATE chunk_sessions
|
||||
SET received_chunks = ?
|
||||
WHERE upload_id = ?
|
||||
""", (json.dumps(received), upload_id))
|
||||
|
||||
log.i(f"Added chunk {chunk_index} to session {upload_id}")
|
||||
return True
|
||||
else:
|
||||
log.w(f"Chunk {chunk_index} already in session {upload_id}")
|
||||
return True
|
||||
except sqlite3.Error as e:
|
||||
log.e(f"Error adding chunk to session {upload_id}: {e}")
|
||||
return False
|
||||
|
||||
def get_received_chunks(upload_id):
|
||||
"""Get list of received chunk indices
|
||||
|
||||
Args:
|
||||
upload_id: 16-character upload ID
|
||||
|
||||
Returns:
|
||||
list: List of chunk indices, or None if session not found
|
||||
"""
|
||||
try:
|
||||
with get_db() as conn:
|
||||
cursor = conn.cursor()
|
||||
row = cursor.execute("""
|
||||
SELECT received_chunks FROM chunk_sessions WHERE upload_id = ?
|
||||
""", (upload_id,)).fetchone()
|
||||
|
||||
if row:
|
||||
return json.loads(row['received_chunks'])
|
||||
return None
|
||||
except sqlite3.Error as e:
|
||||
log.e(f"Error getting received chunks for {upload_id}: {e}")
|
||||
return None
|
||||
|
||||
def delete_session(upload_id):
|
||||
"""Delete chunk upload session
|
||||
|
||||
Args:
|
||||
upload_id: 16-character upload ID
|
||||
|
||||
Returns:
|
||||
bool: True if deleted, False if not found or error
|
||||
"""
|
||||
try:
|
||||
with get_db() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("DELETE FROM chunk_sessions WHERE upload_id = ?", (upload_id,))
|
||||
|
||||
if cursor.rowcount > 0:
|
||||
log.i(f"Deleted chunk session: {upload_id}")
|
||||
return True
|
||||
return False
|
||||
except sqlite3.Error as e:
|
||||
log.e(f"Error deleting chunk session {upload_id}: {e}")
|
||||
return False
|
||||
|
||||
def get_abandoned_sessions(hours=24):
|
||||
"""Get list of abandoned upload sessions older than specified hours
|
||||
|
||||
Args:
|
||||
hours: Age threshold in hours (default: 24)
|
||||
|
||||
Returns:
|
||||
list: List of dicts with session metadata
|
||||
"""
|
||||
try:
|
||||
with get_db() as conn:
|
||||
cursor = conn.cursor()
|
||||
cutoff_time = (datetime.now() - timedelta(hours=hours)).isoformat()
|
||||
|
||||
rows = cursor.execute("""
|
||||
SELECT upload_id, filename, total_chunks, received_chunks, created_at,
|
||||
chunk_dir, max_downloads, expiry_hours
|
||||
FROM chunk_sessions
|
||||
WHERE created_at < ?
|
||||
""", (cutoff_time,)).fetchall()
|
||||
|
||||
sessions = []
|
||||
for row in rows:
|
||||
data = dict(row)
|
||||
data['received_chunks'] = json.loads(data['received_chunks'])
|
||||
sessions.append(data)
|
||||
|
||||
return sessions
|
||||
except sqlite3.Error as e:
|
||||
log.e(f"Error getting abandoned sessions: {e}")
|
||||
return []
|
||||
|
||||
def cleanup_abandoned(hours=24):
|
||||
"""Delete abandoned upload sessions and their chunk directories
|
||||
|
||||
Args:
|
||||
hours: Age threshold in hours (default: 24)
|
||||
|
||||
Returns:
|
||||
int: Number of sessions cleaned up
|
||||
"""
|
||||
try:
|
||||
# Get abandoned sessions first so we can clean up directories
|
||||
abandoned = get_abandoned_sessions(hours)
|
||||
|
||||
if not abandoned:
|
||||
return 0
|
||||
|
||||
with get_db() as conn:
|
||||
cursor = conn.cursor()
|
||||
cutoff_time = (datetime.now() - timedelta(hours=hours)).isoformat()
|
||||
|
||||
# Delete from database
|
||||
cursor.execute("""
|
||||
DELETE FROM chunk_sessions
|
||||
WHERE created_at < ?
|
||||
""", (cutoff_time,))
|
||||
|
||||
count = cursor.rowcount
|
||||
|
||||
# Delete chunk directories
|
||||
for session in abandoned:
|
||||
if os.path.exists(session['chunk_dir']):
|
||||
try:
|
||||
shutil.rmtree(session['chunk_dir'])
|
||||
log.i(f"Deleted chunk directory: {session['chunk_dir']}")
|
||||
except Exception as e:
|
||||
log.w(f"Could not delete chunk directory {session['chunk_dir']}: {e}")
|
||||
|
||||
if count > 0:
|
||||
log.i(f"Cleaned up {count} abandoned upload session(s)")
|
||||
|
||||
return count
|
||||
except sqlite3.Error as e:
|
||||
log.e(f"Error cleaning up abandoned sessions: {e}")
|
||||
return 0
|
||||
|
||||
def update_received_chunks(upload_id, received_chunks):
|
||||
"""Update the entire received_chunks list (used for verification corrections)
|
||||
|
||||
Args:
|
||||
upload_id: 16-character upload ID
|
||||
received_chunks: List of chunk indices
|
||||
|
||||
Returns:
|
||||
bool: True if successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
with get_db() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("""
|
||||
UPDATE chunk_sessions
|
||||
SET received_chunks = ?
|
||||
WHERE upload_id = ?
|
||||
""", (json.dumps(received_chunks), upload_id))
|
||||
|
||||
if cursor.rowcount > 0:
|
||||
log.i(f"Updated received_chunks for session {upload_id}")
|
||||
return True
|
||||
return False
|
||||
except sqlite3.Error as e:
|
||||
log.e(f"Error updating received_chunks for {upload_id}: {e}")
|
||||
return False
|
||||
226
src/models/database.py
Normal file
226
src/models/database.py
Normal file
@@ -0,0 +1,226 @@
|
||||
import sqlite3
|
||||
import json
|
||||
import os
|
||||
from contextlib import contextmanager
|
||||
import utils.mLogger as log
|
||||
|
||||
# Database file location (relative to server directory)
|
||||
DB_DIR = 'db'
|
||||
DB_PATH = os.path.join(DB_DIR, 'file_server.db')
|
||||
|
||||
# Ensure database directory exists
|
||||
if not os.path.exists(DB_DIR):
|
||||
os.makedirs(DB_DIR)
|
||||
|
||||
def get_db_connection():
|
||||
"""Get a thread-safe database connection with row factory"""
|
||||
conn = sqlite3.connect(DB_PATH, check_same_thread=False)
|
||||
conn.row_factory = sqlite3.Row # Access columns by name
|
||||
|
||||
# Enable WAL mode for better concurrency
|
||||
conn.execute("PRAGMA journal_mode=WAL")
|
||||
|
||||
# Set busy timeout to 5 seconds to handle concurrent access
|
||||
conn.execute("PRAGMA busy_timeout=5000")
|
||||
|
||||
return conn
|
||||
|
||||
@contextmanager
|
||||
def get_db():
|
||||
"""Context manager for database connections"""
|
||||
conn = get_db_connection()
|
||||
try:
|
||||
yield conn
|
||||
conn.commit()
|
||||
except Exception as e:
|
||||
conn.rollback()
|
||||
log.e(f"Database error: {e}")
|
||||
raise
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def init_db():
|
||||
"""Initialize database with schema"""
|
||||
log.i("Initializing database...")
|
||||
|
||||
with get_db() as conn:
|
||||
# Admin table
|
||||
conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS admin (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
username TEXT NOT NULL UNIQUE,
|
||||
password_hash TEXT NOT NULL,
|
||||
totp_secret TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL,
|
||||
last_login TEXT
|
||||
)
|
||||
""")
|
||||
|
||||
conn.execute("""
|
||||
CREATE INDEX IF NOT EXISTS idx_admin_username ON admin(username)
|
||||
""")
|
||||
|
||||
# Files table
|
||||
conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS files (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
code TEXT NOT NULL UNIQUE,
|
||||
filename TEXT NOT NULL,
|
||||
path TEXT NOT NULL,
|
||||
max_downloads INTEGER NOT NULL DEFAULT 1,
|
||||
downloads INTEGER NOT NULL DEFAULT 0,
|
||||
expiry_time TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
CONSTRAINT code_length CHECK(length(code) = 6),
|
||||
CONSTRAINT max_downloads_positive CHECK(max_downloads > 0),
|
||||
CONSTRAINT downloads_non_negative CHECK(downloads >= 0)
|
||||
)
|
||||
""")
|
||||
|
||||
conn.execute("""
|
||||
CREATE INDEX IF NOT EXISTS idx_files_code ON files(code)
|
||||
""")
|
||||
|
||||
conn.execute("""
|
||||
CREATE INDEX IF NOT EXISTS idx_files_expiry ON files(expiry_time)
|
||||
""")
|
||||
|
||||
# Chunk sessions table
|
||||
conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS chunk_sessions (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
upload_id TEXT NOT NULL UNIQUE,
|
||||
filename TEXT NOT NULL,
|
||||
total_chunks INTEGER NOT NULL,
|
||||
received_chunks TEXT NOT NULL DEFAULT '[]',
|
||||
created_at TEXT NOT NULL,
|
||||
chunk_dir TEXT NOT NULL,
|
||||
max_downloads INTEGER NOT NULL DEFAULT 1,
|
||||
expiry_hours INTEGER NOT NULL DEFAULT 24,
|
||||
CONSTRAINT upload_id_length CHECK(length(upload_id) = 16),
|
||||
CONSTRAINT total_chunks_positive CHECK(total_chunks > 0)
|
||||
)
|
||||
""")
|
||||
|
||||
conn.execute("""
|
||||
CREATE INDEX IF NOT EXISTS idx_chunk_sessions_upload_id ON chunk_sessions(upload_id)
|
||||
""")
|
||||
|
||||
conn.execute("""
|
||||
CREATE INDEX IF NOT EXISTS idx_chunk_sessions_created_at ON chunk_sessions(created_at)
|
||||
""")
|
||||
|
||||
log.i("Database initialized successfully")
|
||||
|
||||
def migrate_from_json():
|
||||
"""One-time migration from JSON files to SQLite"""
|
||||
|
||||
# Check if we need to migrate
|
||||
admin_json = 'admin_data.json'
|
||||
file_json = 'file_metadata.json'
|
||||
chunk_json = 'chunk_metadata.json'
|
||||
|
||||
# Only migrate if at least one JSON file exists
|
||||
if not any(os.path.exists(f) for f in [admin_json, file_json, chunk_json]):
|
||||
log.i("No JSON files found, skipping migration")
|
||||
return
|
||||
|
||||
log.i("Starting migration from JSON to SQLite...")
|
||||
migrated = False
|
||||
|
||||
with get_db() as conn:
|
||||
# Migrate admin_data.json
|
||||
if os.path.exists(admin_json):
|
||||
try:
|
||||
with open(admin_json, 'r') as f:
|
||||
admin = json.load(f)
|
||||
|
||||
# Check if admin already exists in database
|
||||
existing = conn.execute("SELECT COUNT(*) FROM admin").fetchone()[0]
|
||||
|
||||
if existing == 0:
|
||||
conn.execute("""
|
||||
INSERT INTO admin (username, password_hash, totp_secret, created_at)
|
||||
VALUES (?, ?, ?, ?)
|
||||
""", (admin['username'], admin['password_hash'],
|
||||
admin['totp_secret'], admin['created_at']))
|
||||
log.i(f"Migrated admin user: {admin['username']}")
|
||||
migrated = True
|
||||
else:
|
||||
log.i("Admin already exists in database, skipping admin migration")
|
||||
except Exception as e:
|
||||
log.e(f"Error migrating admin_data.json: {e}")
|
||||
|
||||
# Migrate file_metadata.json
|
||||
if os.path.exists(file_json):
|
||||
try:
|
||||
with open(file_json, 'r') as f:
|
||||
files = json.load(f)
|
||||
|
||||
file_count = 0
|
||||
for code, file_info in files.items():
|
||||
# Check if file already exists
|
||||
existing = conn.execute("SELECT COUNT(*) FROM files WHERE code = ?", (code,)).fetchone()[0]
|
||||
|
||||
if existing == 0:
|
||||
conn.execute("""
|
||||
INSERT INTO files (code, filename, path, max_downloads, downloads, expiry_time)
|
||||
VALUES (?, ?, ?, ?, ?, ?)
|
||||
""", (code, file_info['filename'], file_info['path'],
|
||||
file_info['max_downloads'], file_info['downloads'],
|
||||
file_info['expiry_time']))
|
||||
file_count += 1
|
||||
|
||||
if file_count > 0:
|
||||
log.i(f"Migrated {file_count} file(s)")
|
||||
migrated = True
|
||||
else:
|
||||
log.i("No new files to migrate")
|
||||
except Exception as e:
|
||||
log.e(f"Error migrating file_metadata.json: {e}")
|
||||
|
||||
# Migrate chunk_metadata.json
|
||||
if os.path.exists(chunk_json):
|
||||
try:
|
||||
with open(chunk_json, 'r') as f:
|
||||
chunks = json.load(f)
|
||||
|
||||
chunk_count = 0
|
||||
for upload_id, session in chunks.items():
|
||||
# Check if session already exists
|
||||
existing = conn.execute("SELECT COUNT(*) FROM chunk_sessions WHERE upload_id = ?", (upload_id,)).fetchone()[0]
|
||||
|
||||
if existing == 0:
|
||||
conn.execute("""
|
||||
INSERT INTO chunk_sessions
|
||||
(upload_id, filename, total_chunks, received_chunks, created_at,
|
||||
chunk_dir, max_downloads, expiry_hours)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""", (upload_id, session['filename'], session['total_chunks'],
|
||||
json.dumps(session['received_chunks']), session['created_at'],
|
||||
session['chunk_dir'], session['max_downloads'],
|
||||
session['expiry_hours']))
|
||||
chunk_count += 1
|
||||
|
||||
if chunk_count > 0:
|
||||
log.i(f"Migrated {chunk_count} chunk session(s)")
|
||||
migrated = True
|
||||
else:
|
||||
log.i("No new chunk sessions to migrate")
|
||||
except Exception as e:
|
||||
log.e(f"Error migrating chunk_metadata.json: {e}")
|
||||
|
||||
# Backup JSON files after successful migration
|
||||
if migrated:
|
||||
for filename in [admin_json, file_json, chunk_json]:
|
||||
if os.path.exists(filename):
|
||||
backup_name = f'{filename}.backup'
|
||||
try:
|
||||
os.rename(filename, backup_name)
|
||||
log.i(f"Backed up {filename} to {backup_name}")
|
||||
except Exception as e:
|
||||
log.w(f"Could not backup {filename}: {e}")
|
||||
|
||||
log.i("Migration completed successfully")
|
||||
else:
|
||||
log.i("No migration performed")
|
||||
199
src/models/file_model.py
Normal file
199
src/models/file_model.py
Normal file
@@ -0,0 +1,199 @@
|
||||
import sqlite3
|
||||
import os
|
||||
from datetime import datetime, timedelta
|
||||
from models.database import get_db
|
||||
import utils.mLogger as log
|
||||
|
||||
def create_file(code, filename, path, max_downloads, expiry_hours):
|
||||
"""Create a new file metadata record
|
||||
|
||||
Args:
|
||||
code: 6-character unique code
|
||||
filename: Original filename
|
||||
path: Absolute path to file
|
||||
max_downloads: Maximum number of downloads allowed
|
||||
expiry_hours: Hours until file expires
|
||||
|
||||
Returns:
|
||||
bool: True if successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
with get_db() as conn:
|
||||
expiry_time = datetime.now() + timedelta(hours=expiry_hours)
|
||||
|
||||
conn.execute("""
|
||||
INSERT INTO files (code, filename, path, max_downloads, downloads, expiry_time)
|
||||
VALUES (?, ?, ?, ?, 0, ?)
|
||||
""", (code, filename, path, max_downloads, expiry_time.isoformat()))
|
||||
|
||||
log.i(f"Created file record: {code} - {filename}")
|
||||
return True
|
||||
except sqlite3.Error as e:
|
||||
log.e(f"Error creating file record: {e}")
|
||||
return False
|
||||
|
||||
def get_file(code):
|
||||
"""Get file metadata by code
|
||||
|
||||
Args:
|
||||
code: 6-character file code
|
||||
|
||||
Returns:
|
||||
dict: File metadata with keys: code, filename, path, max_downloads, downloads, expiry_time
|
||||
None: If file not found
|
||||
"""
|
||||
try:
|
||||
with get_db() as conn:
|
||||
cursor = conn.cursor()
|
||||
row = cursor.execute("""
|
||||
SELECT code, filename, path, max_downloads, downloads, expiry_time
|
||||
FROM files
|
||||
WHERE code = ?
|
||||
""", (code,)).fetchone()
|
||||
|
||||
if row:
|
||||
return dict(row)
|
||||
return None
|
||||
except sqlite3.Error as e:
|
||||
log.e(f"Error getting file {code}: {e}")
|
||||
return None
|
||||
|
||||
def increment_download(code):
|
||||
"""Atomically increment download counter for a file
|
||||
|
||||
Args:
|
||||
code: 6-character file code
|
||||
|
||||
Returns:
|
||||
int: New download count, or None if file not found
|
||||
"""
|
||||
try:
|
||||
with get_db() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("""
|
||||
UPDATE files
|
||||
SET downloads = downloads + 1
|
||||
WHERE code = ?
|
||||
""", (code,))
|
||||
|
||||
if cursor.rowcount > 0:
|
||||
# Get the new download count
|
||||
row = cursor.execute("""
|
||||
SELECT downloads FROM files WHERE code = ?
|
||||
""", (code,)).fetchone()
|
||||
|
||||
if row:
|
||||
new_count = row['downloads']
|
||||
log.i(f"Incremented download counter for {code}: {new_count}")
|
||||
return new_count
|
||||
|
||||
return None
|
||||
except sqlite3.Error as e:
|
||||
log.e(f"Error incrementing download for {code}: {e}")
|
||||
return None
|
||||
|
||||
def delete_file(code):
|
||||
"""Delete file metadata record
|
||||
|
||||
Args:
|
||||
code: 6-character file code
|
||||
|
||||
Returns:
|
||||
bool: True if deleted, False if not found or error
|
||||
"""
|
||||
try:
|
||||
with get_db() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("DELETE FROM files WHERE code = ?", (code,))
|
||||
|
||||
if cursor.rowcount > 0:
|
||||
log.i(f"Deleted file record: {code}")
|
||||
return True
|
||||
return False
|
||||
except sqlite3.Error as e:
|
||||
log.e(f"Error deleting file {code}: {e}")
|
||||
return False
|
||||
|
||||
def get_expired_files():
|
||||
"""Get list of expired files or files that exceeded download limit
|
||||
|
||||
Returns:
|
||||
list: List of dicts with file metadata
|
||||
"""
|
||||
try:
|
||||
with get_db() as conn:
|
||||
cursor = conn.cursor()
|
||||
now = datetime.now().isoformat()
|
||||
|
||||
rows = cursor.execute("""
|
||||
SELECT code, filename, path, max_downloads, downloads, expiry_time
|
||||
FROM files
|
||||
WHERE expiry_time < ? OR downloads >= max_downloads
|
||||
""", (now,)).fetchall()
|
||||
|
||||
return [dict(row) for row in rows]
|
||||
except sqlite3.Error as e:
|
||||
log.e(f"Error getting expired files: {e}")
|
||||
return []
|
||||
|
||||
def cleanup_expired():
|
||||
"""Delete expired file records and files that reached download limit
|
||||
|
||||
Returns:
|
||||
int: Number of files cleaned up
|
||||
"""
|
||||
try:
|
||||
# Get expired files first so we can delete physical files
|
||||
expired_files = get_expired_files()
|
||||
|
||||
if not expired_files:
|
||||
return 0
|
||||
|
||||
with get_db() as conn:
|
||||
cursor = conn.cursor()
|
||||
now = datetime.now().isoformat()
|
||||
|
||||
# Delete from database
|
||||
cursor.execute("""
|
||||
DELETE FROM files
|
||||
WHERE expiry_time < ? OR downloads >= max_downloads
|
||||
""", (now,))
|
||||
|
||||
count = cursor.rowcount
|
||||
|
||||
# Delete physical files
|
||||
for file_info in expired_files:
|
||||
if os.path.exists(file_info['path']):
|
||||
try:
|
||||
os.remove(file_info['path'])
|
||||
log.i(f"Deleted expired file: {file_info['path']}")
|
||||
except Exception as e:
|
||||
log.w(f"Could not delete file {file_info['path']}: {e}")
|
||||
|
||||
if count > 0:
|
||||
log.i(f"Cleaned up {count} expired file(s)")
|
||||
|
||||
return count
|
||||
except sqlite3.Error as e:
|
||||
log.e(f"Error cleaning up expired files: {e}")
|
||||
return 0
|
||||
|
||||
def get_all_files():
|
||||
"""Get all file metadata records
|
||||
|
||||
Returns:
|
||||
list: List of dicts with file metadata
|
||||
"""
|
||||
try:
|
||||
with get_db() as conn:
|
||||
cursor = conn.cursor()
|
||||
rows = cursor.execute("""
|
||||
SELECT code, filename, path, max_downloads, downloads, expiry_time, created_at
|
||||
FROM files
|
||||
ORDER BY created_at DESC
|
||||
""").fetchall()
|
||||
|
||||
return [dict(row) for row in rows]
|
||||
except sqlite3.Error as e:
|
||||
log.e(f"Error getting all files: {e}")
|
||||
return []
|
||||
Reference in New Issue
Block a user