nyaa-jelly-tools/nyaa-tools-old/web/server.py

#!/bin/env python3
import os
import re
import secrets
import subprocess
import xml.dom.minidom
from dataclasses import dataclass
from datetime import datetime, timedelta

import requests
import toml
from flask import (
    Flask,
    make_response,
    redirect,
    render_template,
    request,
    url_for,
)

app = Flask(__name__)

CONFIG_DIR = os.environ.get("NYAA_TOOLS_CONFIG_DIR", "/etc/nyaa-tools")
release_groups = toml.load(os.path.join(CONFIG_DIR, "release-groups.toml"))
batch_groups = toml.load(os.path.join(CONFIG_DIR, "batch-groups.toml"))

DEV_ENV = bool(os.environ.get("DEV_ENV", default=False))
if not DEV_ENV:
    PASSWORD = os.environ["PASSWORD"]


active_sessions = {}


def is_token_valid(token):
    return (
        token and token in active_sessions and datetime.now() < active_sessions[token]
    )


# Creates a token and adds it to the "database"
def create_valid_token():
    token = secrets.token_urlsafe(16)

    current_time = datetime.now()
    expiration_time = current_time + timedelta(days=1)

    active_sessions[token] = expiration_time

    return token


def remove_token(token):
    del active_sessions[token]


@app.before_request
def before_request():
    if not DEV_ENV:
        if request.path == "/login":
            return login()
        else:
            if not request.path.startswith("/static"):
                token = request.cookies.get("session_token")
                if not is_token_valid(token):
                    return redirect(url_for("login"))


@dataclass
class SearchItem:
    title: str
    link: str

def search(query) -> list[SearchItem]:
    nyaa_search = requests.get("https://nyaa.si/?page=rss&q={}".format(query)).text
    items = xml.dom.minidom.parseString(nyaa_search).getElementsByTagName("item")

    result = []

    for item in items:
        title: str = item.getElementsByTagName("title")[0].firstChild.data # pyright: ignore[reportOptionalMemberAccess,reportAttributeAccessIssue]
        link: str = item.getElementsByTagName("link")[0].firstChild.data # pyright: ignore[reportOptionalMemberAccess,reportAttributeAccessIssue]

        result.append(SearchItem(title, link))
    return result


@app.route("/")
def index():
    return render_template("index.html")


@app.route("/login", methods=["GET", "POST"])
def login():
    if request.method == "POST":
        password = request.form["password"]

        if password == PASSWORD:
            token = create_valid_token()

            # Set token in a cookie
            resp = make_response(redirect(url_for("index")))
            resp.set_cookie("session_token", token, max_age=60 * 60 * 24)

            return resp
        else:
            return "Invalid credentials. Try again."

    return render_template("login.html")


@app.route("/logout")
def logout():
    # Get the session token
    token = request.cookies.get("session_token")

    remove_token(token)

    # Clear the session and cookie
    resp = make_response(redirect(url_for("login")))
    resp.set_cookie("session_token", "", expires=0)  # Remove cookie
    return resp


@dataclass
class SeriesInfo:
    release_group: str
    title: str
    season: int | None


# Returns the search page if q is empty, else it shows results
@app.route("/series", methods=["GET", "POST"])
def series():
    if request.method == "GET":
        query = request.args.get("q", default="", type=str)

        if query == "":
            return render_template("series/search.html")

        results = search(query)
        parsed_results = []
        for item in results:
            title = item.title
            rg = re.search("\\[([^ ]+)\\].+", title)
            if rg is None or rg.group(1) not in release_groups.keys():
                continue
            info = re.search(release_groups[rg.group(1)], title)
            # TODO: Do not append duplicated entries
            if info is None:
                continue
            try:
                season = int(info.group("season"))
            except IndexError:
                season = None
            parsed_results.append(
                SeriesInfo(
                    release_group=rg.group(1),
                    title=info.group("title"),
                    season=season
                )
            )
        return render_template("series/results.html", results=parsed_results)
    elif request.method == "POST":
        release_group = request.form["release-group"]
        title = request.form["title"]
        season = request.form["season"]
        if release_group and title and season:
            try:
                season = int(season)
            except ValueError:
                return "Invalid season", 400

            series = toml.load(os.path.join(CONFIG_DIR, "series.toml"))
            if "tracking" not in series.keys():
                series["tracking"] = []

            # TODO: Do not allow duplicates
            entry = {"name": title, "release_group": release_group, "season": season}
            series["tracking"].append(entry)
            with open(os.path.join(CONFIG_DIR, "series.toml"), "w") as f:
                toml.dump(series, f)
            return "OK", 200
        else:
            return "Missing parameters", 400


@dataclass
class BatchInfo:
    release_group: str
    title: str
    season: int | None
    episodes: int | None
    link: str

@app.route("/batch", methods=["GET", "POST"])
def batch():
    if request.method == "GET":
        query = request.args.get("q", default="", type=str)

        if query == "":
            return render_template("batch/search.html")

        results = search(query)
        parsed_results = []
        for item in results:
            title = item.title
            rg = re.search("\\[([^ ]+)\\].+", title)
            if rg is None or rg.group(1) not in batch_groups.keys():
                continue
            info = re.search(batch_groups[rg.group(1)], title)
            # TODO: Do not append duplicated entries
            if info is None:
                continue
            try:
                season = int(info.group("season"))
            except IndexError:
                season = None

            try:
                episodes = int(info.group("last")) - int(info.group("first")) + 1
            except IndexError:
                #TODO: Get episodes from file list (requires lxml & cssselector)
                episodes = None

            parsed_results.append(
                BatchInfo(
                    release_group=rg.group(1),
                    title=info.group("title"),
                    season=season,
                    episodes=episodes,
                    link=item.link
                )
            )
        return render_template("batch/results.html", results=parsed_results)
    elif request.method == "POST":
        title = request.form["title"]
        season = request.form["season"]
        link = request.form["link"]
        if title and season and link:
            try:
                season = int(season)
            except ValueError:
                return "Invalid season", 400

            result = subprocess.run(["nyaa-batch", "-t", title, "-s", str(season), "-l", link])
            if result.returncode == 0:
                return "OK", 200
            else:
                return "Error", 500
        else:
            return "Missing parameters", 400

if __name__ == "__main__":
    app.run(debug=DEV_ENV)