monolith/api/views.py

import functools
from json import JSONDecodeError, dumps, loads

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

import main
from modules import userinfo
from utils.logging.log import warn


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("/", methods=["GET"])
    @login_required
    def hello(self, request):
        return "Hello"

    @app.route("/who/<query>/", methods=["POST"])
    @login_required
    def who(self, request, query):
        try:
            data = loads(request.content.read())
        except JSONDecodeError:
            return "Invalid JSON"
        if "query" not in data:
            return "No query provided"
        result = userinfo.getWho(data["query"])
        # Expand the generator
        return dumps({k: [x for x in v] for k, v in result.items()})

    @app.route("/chans/", methods=["POST"])
    @login_required
    def chans(self, request):
        try:
            data = loads(request.content.read())
        except JSONDecodeError:
            return "Invalid JSON"
        if "net" not in data:
            return "No net provided"
        if "query" not in data:
            return "No query provided"
        if not data["query"]:
            warn(f"No query provided: for chans {data}")
            return dumps({})
        result = userinfo.getChansSingle(data["net"], data["query"])
        if not result:
            return dumps({})
        return dumps({"chans": [x for x in result]})

    @app.route("/users/", methods=["POST"])
    @login_required
    def users(self, request):
        try:
            data = loads(request.content.read())
        except JSONDecodeError:
            return "Invalid JSON"
        if "net" not in data:
            return "No net provided"
        if "query" not in data:
            return "No query provided"
        if not data["query"]:
            warn(f"No query provided for users: {data}")
            return dumps({})
        result = userinfo.getUsersSingle(data["net"], data["query"])
        if not result:
            return dumps({})
        return dumps({"users": [x for x in result]})

    @app.route("/online/", methods=["POST"])
    @login_required
    def online(self, request):
        try:
            data = loads(request.content.read())
        except JSONDecodeError:
            return "Invalid JSON"
        if "net" not in data:
            return "No net provided"
        if "query" not in data:
            return "No users provided"
        if not data["query"]:
            warn(f"No users provided: for online {data}")
            return dumps({})
        net = data["net"]
        usermap = {}
        for user in data["query"]:
            channels = userinfo.getChansSingle(net, [user])
            if channels:
                usermap[user] = True
            else:
                usermap[user] = False

        return dumps(usermap)

    @app.route("/num_users/", methods=["POST"])
    @login_required
    def num_users(self, request):
        try:
            data = loads(request.content.read())
        except JSONDecodeError:
            return "Invalid JSON"
        if "net" not in data:
            return "No net provided"
        if "query" not in data:
            return "No users provided"
        if not data["query"]:
            warn(f"No chans provided: for online {data}")
            return dumps({})
        net = data["net"]
        results = userinfo.getUserNum(net, data["query"])

        return dumps(results)

    @app.route("/num_chans/", methods=["POST"])
    @login_required
    def num_chans(self, request):
        try:
            data = loads(request.content.read())
        except JSONDecodeError:
            return "Invalid JSON"
        if "net" not in data:
            return "No net provided"
        if "query" not in data:
            return "No users provided"
        if not data["query"]:
            warn(f"No users provided: for online {data}")
            return dumps({})
        net = data["net"]
        results = userinfo.getChanNum(net, data["query"])

        return dumps(results)