#!/usr/bin/env python3
import re

with open('/root/reservation-system/app/main.py', 'r') as f:
    content = f.read()

# Add auth import
old_import = "from database import get_db, init_db, generate_booking_number, log_change"
new_import = """from database import get_db, init_db, generate_booking_number, log_change
from auth import require_auth, verify_captcha
from login_routes import auth_bp"""
content = content.replace(old_import, new_import)

# Add secret key config after app creation
old_cors = "CORS(app)"
new_cors = """CORS(app)

# Security config
app.secret_key = os.environ.get('SESSION_SECRET', 'dev-secret-change-in-production')

# Register auth blueprint
app.register_blueprint(auth_bp)

@app.before_request
def check_auth():
    # Public endpoints don't require auth
    public_endpoints = ['/', '/api/captcha', '/api/rooms', '/api/availability', 
                       '/api/health', '/api/admin/login']
    if request.path in public_endpoints or request.path.startswith('/static/'):
        return None
    
    # Admin endpoints require login
    if request.path.startswith('/api/admin/') and session.get('user_role') != 'admin':
        return jsonify({"error": "Unauthorized"}), 401"""
content = content.replace(old_cors, new_cors)

# Protect POST reservations with captcha
old_post = '''    if request.method == 'POST':
        data = request.get_json()'''
new_post = '''    if request.method == 'POST':
        data = request.get_json()
        
        # Captcha for non-admin bookings
        if session.get('user_role') != 'admin':
            if not verify_captcha(data.get('captcha_token'), data.get('captcha_answer')):
                return jsonify({"error": "Invalid captcha. Please solve the math problem."}), 403'''
content = content.replace(old_post, new_post, 1)

with open('/root/reservation-system/app/main.py', 'w') as f:
    f.write(content)

print("Security patches applied!")