patch cdn webservice

services/cdn/chall/app.py

@@ -6,6 +6,8 @@ import secrets
 import datetime
 import subprocess
 from pathlib import Path
+from functools import wraps
+
 from flask import *
 from werkzeug.security import generate_password_hash, check_password_hash
 from werkzeug.utils import secure_filename
@@ -13,7 +15,6 @@ from werkzeug.utils import secure_filename
 APP_DIR = os.path.dirname(os.path.abspath(__file__))
 DB_PATH = os.path.join(APP_DIR, "data.db")
 UPLOAD_DIR = os.path.join(APP_DIR, "uploads")
-FLAG_PATH = os.path.join("/flag.txt")
 
 ALLOWED_EXT = {"png", "jpg", "jpeg", "bmp"}
 MAX_CONTENT_LENGTH = 8 * 1024 * 1024
@@ -39,7 +40,8 @@ def close_db(_exc):
 def init_db():
     Path(UPLOAD_DIR).mkdir(parents=True, exist_ok=True)
     db = get_db()
-    db.executescript("""
+    db.executescript(
+        """
     CREATE TABLE IF NOT EXISTS users (
         id INTEGER PRIMARY KEY AUTOINCREMENT,
         username TEXT UNIQUE NOT NULL,
@@ -57,24 +59,12 @@ def init_db():
         created_at TEXT NOT NULL,
         FOREIGN KEY(user_id) REFERENCES users(id)
     );
-    """)
+
+    CREATE INDEX IF NOT EXISTS idx_posts_user_file ON posts(user_id, filename);
+    """
+    )
     db.commit()
 
-def _resolve_flag_file_path():
-    p = FLAG_PATH
-    if os.path.isdir(p):
-        p = os.path.join(p, "flag.txt")
-    os.makedirs(os.path.dirname(p), exist_ok=True)
-    return p
-
-def generate_flag_at_boot():
-    path = _resolve_flag_file_path()
-    if not os.path.exists(path) or os.environ.get("RESEED_FLAG") == "1":
-        token = secrets.token_bytes(32)
-        sha = hashlib.sha256(token).hexdigest()
-        with open(path, "w", encoding="utf-8") as fh:
-            fh.write(f"GEMASTIK{{{sha}}}\n")
-
 def current_user():
     if "uid" not in session:
         return None
@@ -82,6 +72,16 @@ def current_user():
     cur = db.execute("SELECT id, username, role FROM users WHERE id = ?", (session["uid"],))
     return cur.fetchone()
 
+def login_required(view):
+    @wraps(view)
+    def wrapped(*args, **kwargs):
+        if not current_user():
+            dest = request.path
+            flash("Please log in to continue.")
+            return redirect(url_for("login", next=dest))
+        return view(*args, **kwargs)
+    return wrapped
+
 def _is_within(child_path: str, parent_dir: str) -> bool:
     child_real = os.path.realpath(child_path)
     parent_real = os.path.realpath(parent_dir)
@@ -156,8 +156,13 @@ def login():
         if not row or not check_password_hash(row["password_hash"], password):
             flash("Invalid credentials")
             return render_template("login.html")
+
         session["uid"] = row["id"]
         flash(f"Welcome, {row['username']}!")
+
+        next_url = request.args.get("next") or request.form.get("next")
+        if next_url and next_url.startswith("/"):
+            return redirect(next_url)
         return redirect(url_for("gallery"))
     return render_template("login.html")
 
@@ -168,10 +173,9 @@ def logout():
     return redirect(url_for("login"))
 
 @app.route("/upload", methods=["GET", "POST"])
+@login_required
 def upload():
     user = current_user()
-    if not user:
-        return redirect(url_for("login"))
     if request.method == "POST":
         title = request.form.get("title", "").strip() or "(untitled)"
         f = request.files.get("image")
@@ -205,6 +209,7 @@ def upload():
 
 @app.route("/")
 @app.route("/gallery")
+@login_required
 def gallery():
     user = current_user()
     db = get_db()
@@ -212,18 +217,22 @@ def gallery():
         """
         SELECT p.id, p.title, p.filename, p.created_at, u.username
         FROM posts p JOIN users u ON p.user_id = u.id
+        WHERE p.user_id = ?
         ORDER BY p.id DESC
-        """
+        """,
+        (user["id"],),
     )
     posts = cur.fetchall()
     return render_template("gallery.html", posts=posts, user=user)
 
 @app.route("/post/<int:pid>")
+@login_required
 def view_post(pid):
+    user = current_user()
     db = get_db()
     cur = db.execute(
-        "SELECT p.*, u.username FROM posts p JOIN users u ON p.user_id = u.id WHERE p.id = ?",
+        "SELECT p.*, u.username FROM posts p JOIN users u ON p.user_id = u.id WHERE p.id = ? AND p.user_id = ?",
-        (pid,),
+        (pid, user["id"]),
     )
     post = cur.fetchone()
     if not post:
@@ -250,8 +259,20 @@ def view_post(pid):
     return render_template_string(page_src, post=post)
 
 @app.route("/i/<path:filename>")
+@login_required
 def cdn_serve(filename):
-    return send_from_directory(UPLOAD_DIR, filename, as_attachment=False)
+    user = current_user()
+    db = get_db()
+    cur = db.execute(
+        "SELECT 1 FROM posts WHERE user_id = ? AND filename = ? LIMIT 1",
+        (user["id"], filename),
+    )
+    if not cur.fetchone():
+        abort(404)  
+
+    resp = send_from_directory(UPLOAD_DIR, filename, as_attachment=False)
+    resp.headers["Cache-Control"] = "private, max-age=0, no-store"
+    return resp
 
 @app.errorhandler(413)
 def too_large(_):