diff --git a/services/cdn/chall/app.py b/services/cdn/chall/app.py index e2476b9..b603f24 100644 --- a/services/cdn/chall/app.py +++ b/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/") +@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 = ?", - (pid,), + "SELECT p.*, u.username FROM posts p JOIN users u ON p.user_id = u.id WHERE p.id = ? AND p.user_id = ?", + (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/") +@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(_):