From 09196e99b16bf5a0ca4f302819d64adac32036a5 Mon Sep 17 00:00:00 2001 From: edraft Date: Sat, 8 Mar 2025 10:08:03 +0100 Subject: [PATCH] Redirector asgi rewrite --- api/src/redirector.py | 109 +++++++++++++++++------------------------- 1 file changed, 44 insertions(+), 65 deletions(-) diff --git a/api/src/redirector.py b/api/src/redirector.py index 153967f..90d6fa7 100644 --- a/api/src/redirector.py +++ b/api/src/redirector.py @@ -1,9 +1,13 @@ import asyncio import sys -import eventlet -from eventlet import wsgi -from flask import Flask, request, Response, redirect, render_template +import uvicorn +from starlette.applications import Starlette +from starlette.requests import Request +from starlette.responses import RedirectResponse +from starlette.routing import Route, Mount +from starlette.staticfiles import StaticFiles +from starlette.templating import Jinja2Templates from core.database.database import Database from core.environment import Environment @@ -14,72 +18,62 @@ from data.schemas.public.short_url_visit import ShortUrlVisit from data.schemas.public.short_url_visit_dao import shortUrlVisitDao logger = Logger(__name__) +templates = Jinja2Templates(directory="templates") -class Redirector(Flask): - - def __init__(self, *args, **kwargs): - Flask.__init__(self, *args, **kwargs) +async def index(request: Request): + return templates.TemplateResponse("404.html", {"request": request}, status_code=404) -app = Redirector(__name__) - - -@app.route("/") -def index(): - return render_template("404.html"), 404 - - -@app.route("/") -async def _handle_request(path: str): - short_url = await _find_short_url_by_url(path) +async def handle_request(request: Request): + path = request.path_params["path"] + short_url = await shortUrlDao.find_single_by({ShortUrl.short_url: path}) if short_url is None: - return render_template("404.html"), 404 + return templates.TemplateResponse("404.html", {"request": request}, status_code=404) domains = Environment.get("DOMAINS", list[str], []) domain = await short_url.domain logger.debug( - f"Domain: {domain.name if domain is not None else None}, request.host: {request.host}" + f"Domain: {domain.name if domain is not None else None}, request.host: {request.headers['host']}" ) - host = request.host + host = request.headers["host"] if ":" in host: host = host.split(":")[0] domain_strict_mode = Environment.get("DOMAIN_STRICT_MODE", bool, False) if domain is not None and ( - domain.name not in domains - or (domain_strict_mode and not host.endswith(domain.name)) + domain.name not in domains + or (domain_strict_mode and not host.endswith(domain.name)) ): - return render_template("404.html"), 404 + return templates.TemplateResponse("404.html", {"request": request}, status_code=404) user_agent = request.headers.get("User-Agent", "").lower() if "wheregoes" in user_agent or "someothertool" in user_agent: - return await _handle_short_url(path, short_url) + return await _handle_short_url(request, path, short_url) if short_url.loading_screen: - await _track_visit(short_url) + await _track_visit(request, short_url) - return render_template( + return templates.TemplateResponse( "redirect.html", - key=short_url.short_url, - target_url=_get_redirect_url(short_url.target_url), + {"request": request, "key": short_url.short_url, "target_url": _get_redirect_url(short_url.target_url)}, ) - return await _handle_short_url(path, short_url) + return await _handle_short_url(request, path, short_url) -async def _handle_short_url(path: str, short_url: ShortUrl): +async def _handle_short_url(request: Request, path: str, short_url: ShortUrl): if path.startswith("api/"): path = path.replace("api/", "") - await _track_visit(short_url) + await _track_visit(request, short_url) - return _do_redirect(short_url.target_url) + return RedirectResponse(_get_redirect_url(short_url.target_url)) -async def _track_visit(short_url: ShortUrl): +async def _track_visit(request: Request, short_url: ShortUrl): try: await shortUrlVisitDao.create( ShortUrlVisit(0, short_url.id, request.headers.get("User-Agent")) @@ -88,58 +82,43 @@ async def _track_visit(short_url: ShortUrl): logger.error(f"Failed to update short url {short_url.short_url} with error", e) -async def _find_short_url_by_url(url: str) -> ShortUrl: - return await shortUrlDao.find_single_by({ShortUrl.short_url: url}) - - def _get_redirect_url(url: str) -> str: - # todo: multiple protocols like ts3:// if not url.startswith("http://") and not url.startswith("https://"): url = f"http://{url}" return url -def _do_redirect(url: str) -> Response: - return redirect(_get_redirect_url(url)) - - async def configure(): Logger.set_level(Environment.get("LOG_LEVEL", str, "info")) Environment.set_environment(Environment.get("ENVIRONMENT", str, "production")) logger.info(f"Environment: {Environment.get_environment()}") - app.debug = Environment.get_environment() == "development" - await Database.startup_db() +routes = [ + Route("/", endpoint=index), + Mount('/static', StaticFiles(directory='static'), name='static'), + Route("/{path:path}", endpoint=handle_request), +] + +app = Starlette(routes=routes, on_startup=[configure]) + + def main(): if sys.platform == "win32": from asyncio import WindowsSelectorEventLoopPolicy asyncio.set_event_loop_policy(WindowsSelectorEventLoopPolicy()) - loop = asyncio.new_event_loop() - loop.run_until_complete(configure()) - loop.close() - - port = Environment.get("PORT", int, 5001) - logger.info(f"Start API on port: {port}") - if Environment.get_environment() == "development": - logger.info(f"Playground: http://localhost:{port}/") - wsgi.server(eventlet.listen(("0.0.0.0", port)), app, log_output=False) + uvicorn.run( + app, + host="0.0.0.0", + port=Environment.get("PORT", int, 5001), + log_config=None, + ) if __name__ == "__main__": - main() - -# (( -# ( ) -# ; / , -# / \/ -# / | -# / ~/ -# / ) ) ~ edraft -# ___// | / -# --' \_~-, + main() \ No newline at end of file