swingmusic-extended/manage.py

"""
This file is used to run the application.
"""

import os
import logging
from flask_jwt_extended import verify_jwt_in_request
import psutil
import mimetypes
from flask import Response, request

import waitress
import setproctitle

from app.api import create_api
from app.arg_handler import HandleArgs
from app.lib.watchdogg import Watcher as WatchDog
from app.periodic_scan import run_periodic_scans
from app.plugins.register import register_plugins
from app.settings import FLASKVARS, TCOLOR, Keys
from app.setup import run_setup
from app.start_info_logger import log_startup_info
from app.utils.filesystem import get_home_res_path
from app.utils.paths import getClientFilesExtensions
from app.utils.threading import background

mimetypes.add_type("text/css", ".css")

mimetypes.add_type("text/javascript", ".js")
mimetypes.add_type("text/plain", ".txt")
mimetypes.add_type("text/html", ".html")
mimetypes.add_type("image/webp", ".webp")
mimetypes.add_type("image/svg+xml", ".svg")
mimetypes.add_type("image/png", ".png")
mimetypes.add_type("image/vnd.microsoft.icon", ".ico")
mimetypes.add_type("image/gif", ".gif")
mimetypes.add_type("font/woff", ".woff")
mimetypes.add_type("application/manifest+json", ".webmanifest")

werkzeug = logging.getLogger("werkzeug")
werkzeug.setLevel(logging.ERROR)

app = create_api()
app.static_folder = get_home_res_path("client")

# INFO: Routes that don't need authentication
whitelisted_routes = {"/auth/login", "/auth/users", "/auth/logout", "/docs"}
blacklist_extensions = {".webp"}.union(getClientFilesExtensions())


@app.before_request
def verify_auth():
    """
    Verifies the JWT token before each request.
    """
    if request.path == "/" or any(
        request.path.endswith(ext) for ext in blacklist_extensions
    ):
        return

    # if request path starts with any of the blacklisted routes, don't verify jwt
    if any(request.path.startswith(route) for route in whitelisted_routes):
        # print(
        #     "Found whitelisted route: ", request.path, "... Skipping jwt verification"
        # )
        return

    verify_jwt_in_request()


@app.route("/<path:path>")
def serve_client_files(path: str):
    """
    Serves the static files in the client folder.
    """
    js_or_css = path.endswith(".js") or path.endswith(".css")
    if not js_or_css:
        return app.send_static_file(path)

    gzipped_path = path + ".gz"
    user_agent = request.headers.get("User-Agent")

    is_safari = user_agent.find("Safari") >= 0 and user_agent.find("Chrome") < 0

    if is_safari:
        return app.send_static_file(path)

    accepts_gzip = request.headers.get("Accept-Encoding", "").find("gzip") >= 0

    if accepts_gzip:
        if os.path.exists(os.path.join(app.static_folder, gzipped_path)):
            response = app.make_response(app.send_static_file(gzipped_path))
            response.headers["Content-Encoding"] = "gzip"
            return response

    return app.send_static_file(path)


@app.route("/")
def serve_client():
    """
    Serves the index.html file at `client/index.html`.
    """
    return app.send_static_file("index.html")


prev_memory = 0


# INFO: For debugging memory usage
# @app.after_request
def print_memory_usage(response: Response):
    # INFO: Ignore assets
    if (
        request.path.startswith("/img")
        or request.path.endswith(".js")
        or request.path.endswith(".css")
    ):
        return response

    process = psutil.Process(os.getpid())
    global prev_memory
    current_mem = process.memory_info().rss
    diff = (current_mem - prev_memory) / 1024**2
    prev_memory = current_mem

    # INFO: Print memory usage (highlights if diff is more than 0.1 MB)
    print(
        f"\n{request.path} | TOTAL: {current_mem/1024**2} MB | DIFF: {TCOLOR.FAIL if diff > 0.1 else ''}{diff} MB{TCOLOR.ENDC if diff > 0.1 else ''} \n"
    )

    return response


@background
def bg_run_setup() -> None:
    run_periodic_scans()


@background
def start_watchdog():
    WatchDog().run()


@background
def run_swingmusic():
    log_startup_info()
    run_setup()
    bg_run_setup()
    register_plugins()

    start_watchdog()

    setproctitle.setproctitle(f"swingmusic ::{FLASKVARS.get_flask_port()}")


if __name__ == "__main__":
    Keys.load()
    HandleArgs()
    run_swingmusic()

    host = FLASKVARS.get_flask_host()
    port = FLASKVARS.get_flask_port()

    waitress.serve(
        app,
        host=host,
        port=port,
        threads=10,
        ipv6=True,
        ipv4=True,
    )