squirrel/squirrel/main.py
2025-10-03 14:39:30 +10:00

128 lines
3.0 KiB
Python

from flask import (
Flask,
render_template,
redirect,
request,
g,
send_file,
jsonify,
Response,
)
from flask.logging import default_handler
from werkzeug.exceptions import BadRequest, NotFound, Forbidden
import io
from squirrel.conf import settings
from squirrel import auth, repo
from PIL import Image
import os.path
# Enable logging on repo
repo.logger.addHandler(default_handler)
app = Flask(__name__)
app.config["SECRET_KEY"] = "foobar23"
# TODO: better singleton methods
archive = repo.get_archive()
app.register_blueprint(auth.bp)
@app.get("/")
def render_home():
return render_template("home.html")
def search_params(d):
"""
Reference:
"""
try:
params = {
"q": d.get("q"),
"count": int(d.get("count", 20)),
"offset": int(d.get("offset", 0)),
}
except Exception as e:
raise BadRequest(str(e))
if not params["q"]:
raise BadRequest("Missing q parameter")
if params["offset"] < 0:
raise BadRequest("Invalid offset")
return params
def user_repos(user):
if not user:
return []
return user.repos or []
@app.get("/search/<doctype>")
@app.get("/search")
def render_search(doctype="general"):
try:
search = archive.query(**search_params(request.args), extra=user_repos(g.user))
except BadRequest:
return redirect("/")
return render_template(
"search_results.html",
**search,
doctype=doctype,
renderer=f"results_{doctype}.html",
)
@app.get("/search-fragment/<doctype>")
def render_search_fragment(doctype):
search = archive.query(**search_params(request.args), extra=user_repos(g.user))
return render_template(f"results_{doctype}.html", **search)
@app.get("/icon/<type>/<subtype>.svg")
def get_icon(type, subtype):
return send_file(os.path.join(settings["icon_dir"], f"{type}-{subtype}.svg"))
@app.get("/doc/<repo>/<path:udi>")
def get_document(repo, udi):
if repo != "main" and repo not in user_repos(g.user):
raise Forbidden("No access to repo")
filepath = archive.path_for_udi(repo, "/" + udi)
if not os.path.exists(filepath):
raise NotFound("No such document")
return send_file(filepath)
@app.get("/preview/<repo>/<path:docpath>")
def get_preview(repo, docpath):
if repo not in user_repos(g.user):
raise Forbidden("No access to repo")
try:
im = Image.open("/" + docpath)
new_size = (200, 200)
resized_img = im.resize(new_size)
# Save the image to a buffer
img_io = io.BytesIO()
resized_img.save(img_io, "JPEG", quality=70)
img_io.seek(0)
return send_file(img_io, mimetype="image/jpeg")
except Exception as e:
return f"Error: {str(e)}"
@app.get("/api/info")
def get_info():
return jsonify(archive.config)
@app.get("/api/reindex/<repo>")
def get_reindex(repo):
p = archive.reindex(repo)
return Response(p.stderr, mimetype="text/plain")