diff --git a/app.py b/app.py index 0ba18a1..4f2b8d3 100644 --- a/app.py +++ b/app.py @@ -7,7 +7,6 @@ import os app = Flask(__name__) app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///site.db' app.config['SECRET_KEY'] = 'supersecretkey' -app.config['UPLOAD_FOLDER'] = 'uploads' db = SQLAlchemy(app) bcrypt = Bcrypt(app) diff --git a/routes.py b/routes.py index ae77112..8c1ad56 100644 --- a/routes.py +++ b/routes.py @@ -3,8 +3,7 @@ from flask_login import login_user, login_required, logout_user, current_user from app import app, db, bcrypt, login_manager from models import User, Listing, Booking, Image, Document from datetime import datetime -from utils import generate_password, send_email -import os +from utils import generate_password, send_email, save_document @login_manager.user_loader def load_user(user_id): @@ -61,14 +60,17 @@ def upload_document(): if request.method == 'POST': file = request.files['document'] - if file: - filename = file.filename - filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename) - file.save(filepath) - new_document = Document(filename=filename, user_id=current_user.id) + if not file: + return render_template('upload_document.html', message="No file selected.") + + try: + saved_filename = save_document(file, 'documents') + new_document = Document(filename=saved_filename, user_id=current_user.id) db.session.add(new_document) db.session.commit() return redirect(url_for('upload_document')) + except ValueError as e: + return render_template('upload_document.html', message=str(e)) return render_template('upload_document.html', message=None) diff --git a/utils.py b/utils.py index f133295..326d970 100644 --- a/utils.py +++ b/utils.py @@ -2,6 +2,12 @@ import random import string import smtplib from email.mime.text import MIMEText +import os +import uuid +import config +from werkzeug.utils import secure_filename + +UPLOAD_ROOT = os.path.join(os.getcwd(), 'uploads') def generate_password(length=10): characters = string.ascii_letters + string.digits @@ -23,3 +29,22 @@ def send_email(recipient_email, subject, body): server.send_message(msg) except Exception as e: print(f"Failed to send email: {e}") + +def is_file_allowed(file): + return '.' in file.filename and file.filename.rsplit('.', 1)[1].lower() in config.ALLOWED_EXTENSIONS + +def save_document(file, subfolder): + if not is_file_allowed(file): + raise ValueError("File type not allowed. Only images are permitted.") + + if len(file.read()) > config.MAX_FILE_SIZE: + raise ValueError("File size exceeds the limit of 5 MB.") + + file.seek(0) + + upload_folder = os.path.join(UPLOAD_ROOT, subfolder) + os.makedirs(upload_folder, exist_ok=True) + filename = secure_filename(f"{uuid.uuid4().hex}.{file.filename.rsplit('.', 1)[1].lower()}") + filepath = os.path.join(upload_folder, filename) + file.save(filepath) + return filename \ No newline at end of file