mctta/auth/auth.py

from threading import Lock
import sqlite3
import contextlib

class Authentication:
    def __init__(self, database):
        self.db = database
        self.create()

    def connect(self):
        return contextlib.closing(sqlite3.connect(self.db))

    def create(self):
        with self.connect() as c:
            c.execute(
                "CREATE TABLE IF NOT EXISTS users(profileId, token, user)")
            c.execute(
                "CREATE TABLE IF NOT EXISTS new(user)")

    def check_token(self, profileId, token):
        token = token.split(".")[-1]
        if len(token) != 43:
            return False
        with self.connect() as c:
            res = c.execute("SELECT token FROM users WHERE profileId=?",
                (profileId,)).fetchone()
        if res is not None and res[0] != token:
            return False
        return True

    def check_user(self, profileId, token, user):
        with self.connect() as con:
            with con.cursor() as c:
                res = c.execute("SELECT user FROM users WHERE profileId=?",
                    (profileId,)).fetchone()
                if res is not None:
                    return res[0] == user

                # Check if new user can be created
                res = c.execute("SELECT 1 FROM new WHERE user=?",
                    (user,)).fetchone()
                if res is None or res[0] != 1:
                    return False

                # Create new user
                c.execute("DELETE FROM new WHERE user=?", (user,))
                c.execute("INSERT INTO users VALUES (?, ?, ?)",
                    (profileId, token, user))
        return True

    def get_user(self, profileId):
        with self.connect() as c:
            res = c.execute("SELECT user FROM users WHERE profileId=?",
                (profileId)).fetchone()
        if res is None or not isinstance(res[0], str):
            return None
        return res[0]