monolith/api/views.py

import functools
from json import dumps

from klein import Klein
from twisted.web.server import Request

import main
from modules import userinfo


def login_required(func):
    @functools.wraps(func)
    def wrapper(self, *args, **kwargs):
        if isinstance(args[0], Request):
            request = args[0]
            apikey = request.getHeader("ApiKey")
            token = request.getHeader("Token")
            if not apikey:
                return "No API key provided"
            if not token:
                return "No token provided"
            if apikey not in main.tokens:
                return "No such API key"
            config_token = main.tokens[apikey]
            if not token == config_token["hello"]:
                return "Invalid token"
            counter = config_token["counter"]
            request.setHeader("Counter", counter)
        return func(self, *args, **kwargs)

    return wrapper


class API(object):
    """
    Our API webapp.
    """

    app = Klein()

    @app.route("/test/", methods=["GET"])
    @login_required
    def hello(self, request):
        return "Hello"

    @app.route("/who/<query>/", methods=["GET"])
    @login_required
    def who(self, request, query):
        result = userinfo.getWho(query)
        # Expand the generator
        return dumps({k: [x for x in v] for k, v in result.items()})

    @app.route("/chans/<query>/", methods=["GET"])
    @login_required
    def chans(self, request, query):
        query = query.split(",")
        result = userinfo.getChans(query)
        return dumps({k: [x for x in v] for k, v in result.items()})

    @app.route("/users/<query>/", methods=["GET"])
    @login_required
    def users(self, request, query):
        query = [f"#{x}" for x in query.split(",")]
        result = userinfo.getUsers(query)
        return dumps({k: [x for x in v] for k, v in result.items()})