Artberry-web/upload.py

import os
import uuid
import aiofiles
from flask import Blueprint, render_template, redirect, url_for, request, current_app
from flask_login import login_required, current_user
from werkzeug.utils import secure_filename
from models import db, Image, Video, Comic, ComicPage, Post, Cookies, UploadForm, UploadVideoForm
from utils import allowed_file, check_file_content, check_file_size, convert_to_webp, generate_unique_filename

upload_bp = Blueprint('upload', __name__)

@upload_bp.route('/upload', methods=['GET', 'POST'])
@login_required
async def upload():
    form = UploadForm()
    if form.validate_on_submit():
        image_file = form.image_file.data
        tags = form.tags.data
        allowed_mime_types = {'image/png', 'image/jpeg', 'image/gif', 'image/webp'}
        
        if not (allowed_file(image_file.filename, current_app.config['ALLOWED_IMAGE_EXTENSIONS']) and
                check_file_content(image_file, allowed_mime_types) and
                await check_file_size(image_file, current_app.config['MAX_IMAGE_SIZE'])):
            return redirect(url_for('upload.upload'))
        
        unique_filename = f"{uuid.uuid4().hex}.webp"
        filepath = os.path.join(current_app.config['UPLOAD_FOLDER']['images'], unique_filename)
        if os.path.exists(filepath):
            return redirect(url_for('upload.upload'))
        
        webp_image = await convert_to_webp(image_file)
        async with aiofiles.open(filepath, 'wb') as f:
            await f.write(webp_image.read())
        
        img = Image(image_file=unique_filename, username=current_user.username, tags=tags, cookie_votes=0)
        db.session.add(img)
        
        user_cookie = Cookies.query.filter_by(username=current_user.username).first()
        if user_cookie:
            user_cookie.cookies += 1
        else:
            user_cookie = Cookies(username=current_user.username, cookies=1)
            db.session.add(user_cookie)
        
        db.session.commit()
        return redirect(url_for('index'))
    
    return render_template('upload.html', form=form)

@upload_bp.route('/upload_video', methods=['GET', 'POST'])
@login_required
async def upload_video():
    form = UploadVideoForm()
    if form.validate_on_submit():
        video_file = form.video_file.data
        video_thumbnail = form.thumbnail.data
        video_name = form.name.data
        tags = form.tags.data
        description = form.description.data

        allowed_video_mime_types = {'video/mp4', 'video/x-msvideo', 'video/quicktime'}
        allowed_image_mime_types = {'image/png', 'image/jpeg', 'image/gif', 'image/webp'}

        if video_file and video_thumbnail:
            if not (allowed_file(video_file.filename, current_app.config['ALLOWED_VIDEO_EXTENSIONS']) and
                    check_file_content(video_file, allowed_video_mime_types)):
                return redirect(url_for('upload.upload_video'))
            if not await check_file_size(video_file, current_app.config['MAX_VIDEO_SIZE']):
                return redirect(url_for('upload.upload_video'))
            if not (allowed_file(video_thumbnail.filename, current_app.config['ALLOWED_IMAGE_EXTENSIONS']) and
                    check_file_content(video_thumbnail, allowed_image_mime_types)):
                return redirect(url_for('upload.upload_video'))

            video_filename = await generate_unique_filename(current_app.config['UPLOAD_FOLDER']['videos'], 'mp4')
            thumbnail_filename = f"{uuid.uuid4().hex}.webp"

            video_path = os.path.join(current_app.config['UPLOAD_FOLDER']['videos'], video_filename)
            thumbnail_path = os.path.join(current_app.config['UPLOAD_FOLDER']['thumbnails'], thumbnail_filename)

            async with aiofiles.open(video_path, 'wb') as f:
                await f.write(video_file.read())

            webp_thumbnail = await convert_to_webp(video_thumbnail)
            async with aiofiles.open(thumbnail_path, 'wb') as f:
                await f.write(webp_thumbnail.read())

            video = Video(
                video_file=video_filename,
                video_name=video_name,
                video_thumbnail_file=thumbnail_filename,
                username=current_user.username,
                tags=tags,
                description=description
            )
            db.session.add(video)
            db.session.commit()

            return redirect(url_for('videos'))

    return render_template('upload_video.html', form=form)

@upload_bp.route('/comic_upload', methods=['GET', 'POST'])
@login_required
async def comic_upload():
    if request.method == 'POST':
        ct = request.files['thumbnail']
        n = request.form['title']
        tags = request.form.get('tags', '')

        allowed_image_mime_types = {'image/png', 'image/jpeg', 'image/gif', 'image/webp'}

        if db.session.execute(db.select(Comic).filter_by(name=n)).scalar():
            return render_template('comic_upload.html')

        if ct:
            if not (allowed_file(ct.filename, current_app.config['ALLOWED_IMAGE_EXTENSIONS']) and
                    check_file_content(ct, allowed_image_mime_types)):
                return redirect(url_for('upload.comic_upload'))

            tf = f"{uuid.uuid4().hex}.webp"
            tp = os.path.join(current_app.config['UPLOAD_FOLDER']['comicthumbs'], tf)
            webp_thumbnail = await convert_to_webp(ct)
            async with aiofiles.open(tp, 'wb') as f:
                await f.write(webp_thumbnail.read())

            cf = os.path.join(current_app.config['UPLOAD_FOLDER']['comics'], n)
            os.makedirs(cf, exist_ok=True)

            new_comic = Comic(
                comic_folder=n,
                comic_thumbnail_file=tf,
                username=current_user.username,
                name=n,
                tags=tags
            )
            db.session.add(new_comic)
            db.session.flush()

            async def save_pages():
                pages = request.files.getlist('pages[]')
                for i, p in enumerate(sorted(pages, key=lambda x: x.filename), start=1):
                    if p:
                        if not (allowed_file(p.filename, current_app.config['ALLOWED_IMAGE_EXTENSIONS']) and
                                check_file_content(p, allowed_image_mime_types)):
                            return redirect(url_for('upload.comic_upload'))

                        filename = f"{uuid.uuid4().hex}.webp"
                        file_path = os.path.join(cf, filename)
                        webp_page = await convert_to_webp(p)
                        async with aiofiles.open(file_path, 'wb') as f:
                            await f.write(webp_page.read())
                        
                        new_page = ComicPage(comic_id=new_comic.id, page_number=i, file_path=file_path)
                        db.session.add(new_page)
                db.session.commit()
            
            await save_pages()
            return redirect(url_for('comics'))
    
    return render_template('comic_upload.html')

@upload_bp.route('/upload_post', methods=['GET', 'POST'])
@login_required
async def upload_post():
    if request.method == 'POST':
        post_text = request.form.get('post_text')
        post_media = request.files.get('post_media')

        if post_text:
            new_post = Post(
                username=current_user.username,
                text=post_text
            )
            db.session.add(new_post)

            if post_media and allowed_file(post_media.filename, current_app.config['ALLOWED_IMAGE_EXTENSIONS']):
                if await check_file_size(post_media, current_app.config['MAX_IMAGE_SIZE']):
                    unique_filename = f"{uuid.uuid4().hex}.webp"
                    media_path = os.path.join(current_app.config['UPLOAD_FOLDER']['posts'], unique_filename)
                    webp_image = await convert_to_webp(post_media)
                    async with aiofiles.open(media_path, 'wb') as f:
                        await f.write(webp_image.read())
                    new_post.media_file = unique_filename
                else:
                    return redirect(url_for('upload.upload_post'))

            db.session.commit()
            return redirect(url_for('user_posts', username=current_user.username))
        else:
            return redirect(url_for('upload.upload_post'))

    return render_template('upload_post.html')