diff --git a/src/app.py b/src/app.py index ab77985..d1be45c 100644 --- a/src/app.py +++ b/src/app.py @@ -14,22 +14,30 @@ from controllers.admin_controller import ( update_config, regenerate_key ) from views import views -from models import init_db, migrate_from_json, migrate_client_key_to_db +from models import init_db, migrate_from_json, migrate_client_key_to_db, get_secret_key -app = Flask(__name__) -app.config['UPLOAD_FOLDER'] = 'tmp/uploads' -app.config['SECRET_KEY'] = os.environ.get('SECRET_KEY', os.urandom(24).hex()) -app.config['PERMANENT_SESSION_LIFETIME'] = timedelta(hours=24) - -if not os.path.exists(app.config['UPLOAD_FOLDER']): - os.makedirs(app.config['UPLOAD_FOLDER']) - -# Initialize database +# Initialize database first (before creating Flask app) log.i("Initializing database...") init_db() migrate_from_json() migrate_client_key_to_db() +# Get secret key from database +secret_key = get_secret_key() +if not secret_key: + log.e("Failed to get secret key from database!") + secret_key = os.urandom(24).hex() # Fallback only +else: + log.i(f"Loaded secret key from database (length: {len(secret_key)})") + +app = Flask(__name__) +app.config['UPLOAD_FOLDER'] = 'tmp/uploads' +app.config['SECRET_KEY'] = secret_key +app.config['PERMANENT_SESSION_LIFETIME'] = timedelta(hours=24) + +if not os.path.exists(app.config['UPLOAD_FOLDER']): + os.makedirs(app.config['UPLOAD_FOLDER']) + # Run cleanup on startup cleanup_expired_files() cleanup_abandoned_uploads() diff --git a/src/controllers/admin_controller.py b/src/controllers/admin_controller.py index 1d073c1..de68565 100644 --- a/src/controllers/admin_controller.py +++ b/src/controllers/admin_controller.py @@ -89,7 +89,7 @@ def setup_admin_totp(): # Convert QR code to base64 buffer = io.BytesIO() - img.save(buffer, format='PNG') + img.save(buffer) # PyPNG doesn't use format parameter qr_base64 = base64.b64encode(buffer.getvalue()).decode() return render_template('admin_setup_totp.html', @@ -113,7 +113,7 @@ def setup_admin_totp(): qr.make(fit=True) img = qr.make_image(fill_color="black", back_color="white") buffer = io.BytesIO() - img.save(buffer, format='PNG') + img.save(buffer) qr_base64 = base64.b64encode(buffer.getvalue()).decode() return render_template('admin_setup_totp.html', @@ -137,7 +137,7 @@ def setup_admin_totp(): qr.make(fit=True) img = qr.make_image(fill_color="black", back_color="white") buffer = io.BytesIO() - img.save(buffer, format='PNG') + img.save(buffer) qr_base64 = base64.b64encode(buffer.getvalue()).decode() return render_template('admin_setup_totp.html', @@ -287,7 +287,7 @@ def generate_config_qr(server_address, server_port, client_key): # Convert to base64 buffer = io.BytesIO() - img.save(buffer, format='PNG') + img.save(buffer) qr_base64 = base64.b64encode(buffer.getvalue()).decode() return qr_base64 diff --git a/src/models/__init__.py b/src/models/__init__.py index 08757d3..3cb9aab 100644 --- a/src/models/__init__.py +++ b/src/models/__init__.py @@ -38,6 +38,7 @@ from .chunk_model import ( from .config_model import ( get_config, get_client_key, + get_secret_key, init_config, update_server_address, update_server_port, diff --git a/src/models/config_model.py b/src/models/config_model.py index aa26f52..5c28cfe 100644 --- a/src/models/config_model.py +++ b/src/models/config_model.py @@ -10,18 +10,23 @@ def generate_client_key(): return secrets.token_hex(32) # 32 bytes = 64 hex chars +def generate_secret_key(): + """Generate a random 64-character hex string for Flask SECRET_KEY""" + return secrets.token_hex(32) # 32 bytes = 64 hex chars + + def get_config(): """Get server configuration (singleton) Returns: - dict: Configuration with client_key, server_address, server_port, timestamps + dict: Configuration with client_key, secret_key, server_address, server_port, timestamps None: If config doesn't exist or error occurred """ try: with get_db() as conn: cursor = conn.cursor() row = cursor.execute(""" - SELECT client_key, server_address, server_port, created_at, updated_at + SELECT client_key, secret_key, server_address, server_port, created_at, updated_at FROM server_config WHERE id = 1 """).fetchone() @@ -45,11 +50,23 @@ def get_client_key(): return config['client_key'] if config else None -def init_config(client_key=None): - """Initialize server config with provided or generated key +def get_secret_key(): + """Get the Flask secret key from database + + Returns: + str: The secret key + None: If config doesn't exist or error occurred + """ + config = get_config() + return config['secret_key'] if config else None + + +def init_config(client_key=None, secret_key=None): + """Initialize server config with provided or generated keys Args: client_key (str, optional): Client key to use. If None, generates new key + secret_key (str, optional): Secret key to use. If None, generates new key Returns: bool: True if successful, False otherwise @@ -63,20 +80,26 @@ def init_config(client_key=None): log.i("Server config already exists, skipping initialization") return True - # Generate key if not provided + # Generate keys if not provided if not client_key: client_key = generate_client_key() log.i("Generated new client key") else: log.i("Using provided client key from environment") + if not secret_key: + secret_key = generate_secret_key() + log.i("Generated new secret key") + else: + log.i("Using provided secret key from environment") + # Insert initial config now = datetime.now().isoformat() conn.execute(""" - INSERT INTO server_config (id, client_key, created_at, updated_at) - VALUES (1, ?, ?, ?) - """, (client_key, now, now)) + INSERT INTO server_config (id, client_key, secret_key, created_at, updated_at) + VALUES (1, ?, ?, ?, ?) + """, (client_key, secret_key, now, now)) log.i("Server config initialized successfully") return True diff --git a/src/models/database.py b/src/models/database.py index 80e4e3b..2078a24 100644 --- a/src/models/database.py +++ b/src/models/database.py @@ -9,8 +9,7 @@ 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) +os.makedirs(DB_DIR, exist_ok=True) def get_db_connection(): """Get a thread-safe database connection with row factory""" @@ -115,11 +114,13 @@ def init_db(): CREATE TABLE IF NOT EXISTS server_config ( id INTEGER PRIMARY KEY CHECK(id = 1), client_key TEXT NOT NULL, + secret_key TEXT, server_address TEXT, server_port INTEGER, created_at TEXT NOT NULL DEFAULT (datetime('now')), updated_at TEXT NOT NULL DEFAULT (datetime('now')), - CONSTRAINT client_key_length CHECK(length(client_key) = 64) + CONSTRAINT client_key_length CHECK(length(client_key) = 64), + CONSTRAINT secret_key_length CHECK(secret_key IS NULL OR length(secret_key) = 64) ) """) @@ -249,6 +250,8 @@ def migrate_client_key_to_db(): if existing > 0: log.i("Server config already exists in database, skipping migration") + # Check if secret_key needs to be added to existing config + migrate_secret_key() return # Get CLIENT_KEY from environment if set @@ -261,3 +264,29 @@ def migrate_client_key_to_db(): else: log.i("No valid CLIENT_KEY in environment, generating new key") init_config() + + +def migrate_secret_key(): + """Migrate existing server_config to add secret_key if missing""" + from models.config_model import generate_secret_key + + try: + with get_db() as conn: + # Check if secret_key already exists + result = conn.execute("SELECT secret_key FROM server_config WHERE id = 1").fetchone() + + if result and result[0]: + log.i("Secret key already exists in database") + return + + # Generate and add secret_key + secret_key = generate_secret_key() + conn.execute(""" + UPDATE server_config + SET secret_key = ? + WHERE id = 1 + """, (secret_key,)) + + log.i("Added secret key to existing server config") + except Exception as e: + log.e(f"Error migrating secret key: {e}")