uotp/uotp/helper.py

import pyotp
import requests

from os.path import join
from .config import env_vars


def is_valid_otp(etcd_client, name, realm, token):
    _key = join(env_vars.get("BASE_PREFIX"), name)
    entry = etcd_client.get(_key, value_in_json=True)
    if entry:
        if realm not in entry.value["realm"]:
            return False

        totp = pyotp.TOTP(entry.value["seed"])
        try:
            is_token_valid = totp.verify(token)
        except Exception:
            return False
        else:
            return is_token_valid

    return False


def create_admin_if_dont_exists(etcd_client):
    _key = join(env_vars.get("BASE_PREFIX"), "admin")
    if etcd_client.get(_key) is None:
        print("admin does not exists!. So, creating one")
        _value = {
            "seed": pyotp.random_base32(),
            "realm": [env_vars.get("ADMIN_REALM")],
        }
        etcd_client.put(_key, _value, value_in_json=True)