add badges

app/routes/users.py

@@ -8,7 +8,7 @@ from ..models import (
     Users, Sessions, Subscriptions,
     Avatars, PasswordResetLinks, InviteKeys,
     BookmarkCollections, BookmarkedThreads,
-    Mentions, PostHistory,
+    Mentions, PostHistory, Badges, BadgeUploads,
     )
 from ..constants import InfoboxKind, PermissionLevel, SIG_BANNED_TAGS
 from ..auth import digest, verify
@@ -21,6 +21,7 @@ import re
 import os
 
 AVATAR_MAX_SIZE = 1000 * 1000
+BADGE_MAX_SIZE = 1000 * 500
 
 bp = Blueprint("users", __name__, url_prefix = "/users/")
 
@@ -56,6 +57,21 @@ def validate_and_create_avatar(input_image, filename):
     except WandException:
         return False
 
+def validate_and_create_badge(input_image, filename):
+    try:
+        with Image(blob=input_image) as img:
+            if img.width != 88 or img.height != 31:
+                return False
+            if hasattr(img, 'sequence') and len(img.sequence) > 1:
+                img = Image(image=img.sequence[0])
+            img.strip()
+
+            img.format = 'webp'
+            img.compression_quality = 90
+            img.save(filename=filename)
+            return True
+    except WandException:
+        return False
 
 def is_logged_in():
     return "pyrom_session_key" in session
@@ -855,3 +871,96 @@ def delete_page_confirm(username):
     session.clear()
     target_user.delete()
     return redirect(url_for('topics.all_topics'))
+
+
+@bp.post('/<username>/save-badges')
+@login_required
+@redirect_to_own
+def save_badges(username):
+    user = get_active_user()
+    badge_choices = request.form.getlist('badge_choice[]')
+    badge_files = request.files.getlist('badge_file[]')
+    badge_labels = request.form.getlist('badge_label[]')
+    badge_links = request.form.getlist('badge_link[]')
+
+    if not (len(badge_choices) == len(badge_files) == len(badge_labels) == len(badge_links)):
+        return 'nope'
+    pending_badges = []
+    rejected_filenames = []
+
+    # lack of file can be checked with a simple `if not file`
+
+    for i in range(len(badge_choices)):
+        is_custom = badge_choices[i] == 'custom'
+        file = badge_files[i]
+        pending_badge = {
+            'upload': badge_choices[i],
+            'is_custom': is_custom,
+            'label': badge_labels[i],
+            'link': badge_links[i],
+            'sort_order': i,
+        }
+        if is_custom:
+            file.seek(0, os.SEEK_END)
+            file_size = file.tell()
+            file.seek(0, os.SEEK_SET)
+
+            if file_size >= BADGE_MAX_SIZE:
+                rejected_filenames.append(file.filename)
+                continue
+
+            file_bytes = file.read()
+
+            pending_badge['original_filename'] = file.filename
+
+            now = int(time.time())
+            filename = f'u{user.id}d{now}s{i}.webp'
+            output_path = os.path.join(current_app.config['BADGES_UPLOAD_PATH'], filename)
+            proxied_filename = f'/static/badges/user/{filename}'
+            res = validate_and_create_badge(file_bytes, output_path)
+            if not res:
+                rejected_filenames.append(file.filename)
+                continue
+
+            pending_badge['proxied_filename'] = proxied_filename
+            pending_badge['uploaded_at'] = now
+
+        pending_badges.append(pending_badge)
+
+    if rejected_filenames:
+        flash(f'Invalid badges.;Some of your uploaded badges are incorrect: {", ".join(rejected_filenames)}. Your badges have not been modified.', InfoboxKind.ERROR)
+        return redirect(url_for('.settings', username=user.username))
+
+    with db.transaction():
+        existing_badges = Badges.findall({'user_id': int(user.id)})
+        for badge in existing_badges:
+            badge.delete()
+
+    with db.transaction():
+        for pending_badge in pending_badges:
+            if pending_badge['is_custom']:
+                bu = BadgeUploads.create({
+                    'file_path': pending_badge['proxied_filename'],
+                    'uploaded_at': pending_badge['uploaded_at'],
+                    'original_filename': pending_badge['original_filename'],
+                    'user_id': int(user.id),
+                })
+            else:
+                bu = BadgeUploads.find({
+                    'id': int(pending_badge['upload'])
+                })
+            badge = Badges.create({
+                'user_id': int(user.id),
+                'upload': int(bu.id),
+                'label': pending_badge['label'],
+                'link': pending_badge['link'],
+                'sort_order': pending_badge['sort_order']
+            })
+
+        for stale_upload in BadgeUploads.get_unused_for_user(user.id):
+            filename = os.path.join(current_app.config['BADGES_UPLOAD_PATH'], os.path.basename(stale_upload.file_path))
+            os.remove(filename)
+            stale_upload.delete()
+
+    flash('Badges saved.', InfoboxKind.INFO)
+    return redirect(url_for('.settings', username=user.username))