MsRewards/modules/Classes/UserCredentials.py

from pyotp import TOTP

from modules.Tools.logger import debug, warning


class UserCredentials:
    def __init__(self):
        self.data = {}
        self.current = 0
        self.total = 0

    def add(self, username: str, password: str, tfa: str = None):
        debug(
            f"adding account with data : Username: {username}, Password: {password}, 2FA: {'None' if tfa == '' else tfa}")
        self.data[self.total] = {
            "username": username,
            "password": password,
            "2fa": None if tfa == '' else tfa
        }
        self.total += 1

    def tfa_enable(self):
        return self.data[self.current]["2fa"] is not None

    def get_mail(self):
        return self.data[self.current]["username"]

    def get_password(self):
        return self.data[self.current]["password"]

    def get_tfa(self):
        if not self.tfa_enable():
            warning("Warning: TFA is not enabled. Can't get a TFA code.")
            return None
        return TOTP(self.data[self.current]["2fa"])

    def next_account(self):
        self.current += 1
        if self.is_valid():
            debug(f"New credentials: {self.data[self.current]}")
        else:
            debug("No new credentials.")

    def is_valid(self):
        return (self.current < self.total
                and self.get_mail() != "" and self.get_mail is not None)