app.py

import json
import time

import jwt
import requests
from flask import Flask, request, redirect, g, session, url_for, render_template
from jwt import PyJWKClient
from peewee import SqliteDatabase, Model, CharField, TextField, IntegrityError

import config

app = Flask(__name__)
app.config.from_object(config)

jwk_client = PyJWKClient("https://usetrmnl.com/.well-known/jwks.json")

database = SqliteDatabase(app.config["DATABASE"])


class BaseModel(Model):
    class Meta:
        database = database


class PluginInstall(BaseModel):
    plugin_setting_id = CharField(null=True)
    plugin_uuid = CharField(null=True)
    access_token = CharField(null=True)
    data = TextField(null=True)


@app.route("/new")
def new():
    # new plugin install endpoint
    req_body = {
        "client_id": app.config["CLIENT_ID"],
        "client_secret": app.config["CLIENT_SECRET"],
        "grant_type": "authorization_code",
        "code": request.args["code"],
    }

    response = requests.post("https://usetrmnl.com/oauth/token", data=req_body)
    try:
        with database.atomic():
            pi = PluginInstall.create(
                access_token=response.json()["access_token"],
            )
    except IntegrityError:
        pass  # already installed, redirect

    return redirect(request.args["installation_callback_url"])


@app.route("/install", methods=["POST"])
def install_webhook():
    # plugin installed webhook
    access_token = request.headers["Authorization"].split(" ")[1]
    data = request.get_json()
    # get pi
    pi = PluginInstall.get_or_none(PluginInstall.access_token == access_token)
    if not pi:
        return "", 404

    data["current_task"] = "Figuring out life.."

    pi.data = json.dumps(data)
    pi.plugin_setting_id = data["user"]["plugin_setting_id"]
    pi.plugin_uuid = data["user"]["uuid"]
    pi.save()

    return "OK"


@app.route("/settings/login")
def settings_login():
    # management endpoint
    token = request.args["jwt"]
    public_key = jwk_client.get_signing_key_from_jwt(token).key
    try:
        decoded = jwt.decode(
            token, public_key, algorithms=["RS256"], audience=app.config["CLIENT_ID"]
        )
    except Exception as e:
        print(e)
        return "JWT failed to decode", 403

    pi = PluginInstall.get_or_none(PluginInstall.plugin_uuid == decoded["sub"])

    if not pi:
        return "Failed to find plugin settings.", 404

    session["pi"] = pi.id
    session["expires"] = int(time.time()) + (60 * 30)

    return redirect(url_for("settings"))


def logout_and_redirect():
    session.clear()
    return redirect("https://usetrmnl.com")


@app.route("/settings", methods=["GET", "POST"])
def settings():
    # this is not the route we give to TRMNL management URL, but a second route we use internally
    # to clear the token and uuid from the browser logs
    # we use cookies to handle this "login" feature
    if "pi" not in session or "expires" not in session:
        return logout_and_redirect()

    if session["expires"] < int(time.time()):
        return logout_and_redirect()

    pi = PluginInstall.get_or_none(session["pi"])
    if not pi:
        return logout_and_redirect()

    data = json.loads(pi.data)["user"]

    # Get the current task from the data or set a default value
    pi_data = json.loads(pi.data)
    current_task = pi_data.get("current_task", "")

    # Handle form submission
    if request.method == "POST":
        current_task = request.form.get("current_task", "")

        # Update the data with the new current_task
        pi_data["current_task"] = current_task
        pi.data = json.dumps(pi_data)
        pi.save()

        # log the user out and redirect them back to TRMNL
        session.clear()
        return redirect(
            f"https://usetrmnl.com/plugin_settings/{pi.plugin_setting_id}/edit?force_refresh=true"
        )

    # Render the template with the data
    return render_template(
        "settings.j2",
        first_name=data["first_name"],
        plugin_uuid=pi.plugin_uuid,
        plugin_setting_id=pi.plugin_setting_id,
        current_task=current_task,
    )


@app.route("/markup", methods=["POST"])
def markup():
    # markup endpoint
    access_token = request.headers["Authorization"].split(" ")[1]
    pi = PluginInstall.get_or_none(PluginInstall.access_token == access_token)
    if not pi:
        return "Failed to find plugin settings. ", 404

    data = json.loads(pi.data)
    current_task = data.get("current_task", "")

    ret = {
        "markup": render_template(
            "markup.j2", current_task=current_task, view_type="full"
        )
    }

    for view_type in [
        "half_horizontal",
        "half_vertical",
        "quadrant",
    ]:
        ret[f"markup_{view_type}"] = render_template(
            "markup.j2", current_task=current_task, view_type=view_type
        )
    return ret


@app.route("/uninstall", methods=["POST"])
def uninstall():
    # uninstall endpoint
    access_token = request.headers["Authorization"].split(" ")[1]
    pi = PluginInstall.get_or_none(PluginInstall.access_token == access_token)
    if not pi:
        return "", 404

    pi.delete_instance()
    return "OK"


@app.before_request
def before_request():
    g.db = database
    g.db.connect()


@app.after_request
def after_request(response):
    g.db.close()
    return response


with database:
    database.create_tables([PluginInstall])

if __name__ == "__main__":
    app.run(debug=True, port=8009)