PYTHON

Secure File Uploads: Type and Size Validation

Implement secure file upload handling in Python Flask by robustly validating file extensions, MIME types, and sizes server-side to prevent malicious uploads and resource exhaustion attacks.

import os
from flask import Flask, request, redirect, url_for, flash
from werkzeug.utils import secure_filename

# Ensure to install Flask and Werkzeug (pip install Flask)
# The `secure_filename` function is part of Werkzeug, which Flask depends on.

app = Flask(__name__)
app.secret_key = 'supersecretkey' # Used for flashing messages, crucial for production
app.config['UPLOAD_FOLDER'] = 'uploads' # Store files outside web root if possible
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024 # 16 MB max upload size

# Define allowed extensions and corresponding MIME types
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'pdf'}
ALLOWED_MIMETYPES = {'image/png', 'image/jpeg', 'image/gif', 'application/pdf'}

def allowed_file(filename, mimetype):
    """Checks if the filename has an allowed extension and a matching MIME type."""
    ext = '.' in filename and filename.rsplit('.', 1)[1].lower()
    if ext not in ALLOWED_EXTENSIONS:
        return False, "Invalid file extension."

    # Perform basic MIME type check (server-side check is more reliable)
    # Note: Client-side MIME type can be spoofed.
    if mimetype not in ALLOWED_MIMETYPES:
        return False, "Invalid file MIME type."

    # Example of a more robust MIME type check (requires 'python-magic' or similar)
    # import magic
    # try:
    #     actual_mime = magic.from_buffer(request.files['file'].read(1024), mime=True)
    #     request.files['file'].seek(0) # Reset stream position
    #     if actual_mime not in ALLOWED_MIMETYPES:
    #         return False, "Actual file content MIME type not allowed."
    # except Exception as e:
    #     print(f"Error checking actual MIME type: {e}")
    #     return False, "Error processing file content."
    
    return True, ""

@app.route('/upload', methods=['GET', 'POST'])
def upload_file():
    if request.method == 'POST':
        # Check if the post request has the file part
        if 'file' not in request.files:
            flash('No file part')
            return redirect(request.url)
        file = request.files['file']
        # If user does not select file, browser also submits an empty part without filename
        if file.filename == '':
            flash('No selected file')
            return redirect(request.url)

        is_allowed, msg = allowed_file(file.filename, file.mimetype)
        if file and is_allowed:
            filename = secure_filename(file.filename)
            file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
            os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True) # Ensure folder exists
            file.save(file_path)
            flash(f'File {filename} successfully uploaded.')
            return redirect(url_for('upload_file'))
        else:
            flash(f'Upload failed: {msg or "File type not allowed."}')
            return redirect(request.url)

    return '''
<!doctype html>
<title>Upload new File</title>
<h1>Upload new File</h1>
{% with messages = get_flashed_messages() %}
  {% if messages %}
    <ul class=flashes>
    {% for message in messages %}
      <li>{{ message }}</li>
    {% endfor %}
    </ul>
  {% endif %}
{% endwith %}
<form method=post enctype=multipart/form-data>
  <input type=file name=file>
  <input type=submit value=Upload>
</form>
'''

if __name__ == '__main__':
    app.run(debug=True) # Set debug=False in production
How it works: This Python Flask snippet demonstrates secure handling of file uploads by implementing server-side validation for file type and size. It defines `ALLOWED_EXTENSIONS` and `ALLOWED_MIMETYPES` as whitelists to ensure only specific, safe file types (like images and PDFs) can be uploaded. The `allowed_file` function checks both the file's extension and its reported MIME type. While client-provided MIME types can be spoofed, this initial check is a good first line of defense. For production, more robust MIME type detection (e.g., using `python-magic`) is recommended, as shown in the commented-out section. Furthermore, `app.config['MAX_CONTENT_LENGTH']` limits the upload size to prevent denial-of-service attacks, and `secure_filename()` sanitizes filenames to prevent directory traversal vulnerabilities. Files are saved to a designated `UPLOAD_FOLDER`, ideally outside the web-accessible root. This multi-layered approach is crucial for preventing malicious file execution or resource exhaustion on your server.

Need help integrating this into your project?

Our team of expert developers can help you build your custom application from scratch.

Hire DigitalCodeLabs