codekingpro/portable-devtools

1"""2    flask_security.totp3    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~4 5    Flask-Security TOTP (Timed-One-Time-Passwords) module6 7    :copyright: (c) 2019-2024 by J. Christopher Wagner (jwag).8    :license: MIT, see LICENSE for more details.9"""10 11from __future__ import annotations12 13import base6414import io15import typing as t16 17from passlib.totp import TOTP, TokenError, TotpMatch18from passlib.pwd import genword19 20if t.TYPE_CHECKING:  # pragma: no cover21    from .datastore import User22 23 24class Totp:25    """Encapsulate usage of Passlib TOTP functionality.26 27    Flask-Security doesn't implement any replay-attack protection out of the box28    as suggested by:29    https://passlib.readthedocs.io/en/stable/narr/totp-tutorial.html#match-verify30 31    Subclass this and implement the get/set last_counter methods. Your subclass can32    be registered at Flask-Security creation/initialization time.33 34    .. versionadded:: 3.4.035 36    """37 38    def __init__(self, secrets: dict[str | int, str], issuer: str):39        """Initialize a totp factory.40        secrets are used to encrypt the per-user totp_secret on disk.41        """42        # This should be a dict with at least one entry43        if not isinstance(secrets, dict) or len(secrets) < 1:44            raise ValueError("secrets needs to be a dict with at least one entry")45        self._totp = TOTP.using(issuer=issuer, secrets=secrets)46 47    def generate_totp_password(self, totp_secret: str) -> str:48        """Get time-based one-time password on the basis of given secret and time49        :param totp_secret: the unique shared secret of the user50        """51        return self._totp.from_source(totp_secret).generate().token52 53    def generate_totp_secret(self) -> str:54        """Create new user-unique totp_secret.55 56        We return an encrypted json string so that when sent in a cookie or57        sent to DB - it is encrypted.58 59        """60        return self._totp.new().to_json(encrypt=True)61 62    def verify_totp(63        self, token: str, totp_secret: str, user: User, window: int = 064    ) -> bool:65        """Verifies token for specific user.66 67        :param token: token to be check against user's secret68        :param totp_secret: the unique shared secret of the user69        :param user: User model70        :param window: optional. How far backward and forward in time to search71         for a match. Measured in seconds.72        :return: True if match73        """74 75        # TODO - in old implementation  using onetimepass window was described76        # as 'compensate for clock skew) and 'interval_length' would say how long77        # the token is good for.78        # In passlib - 'window' means how far back and forward to look and 'clock_skew'79        # is specifically for well, clock slew.80        try:81            tmatch = self._totp.verify(82                token,83                totp_secret,84                window=window,85                last_counter=self.get_last_counter(user),86            )87            self.set_last_counter(user, tmatch)88            return True89 90        except TokenError:91            return False92 93    def get_totp_uri(self, username: str, totp_secret: str) -> str:94        """Generate provisioning url for use with the qrcode95                scanner built into the app96 97        :param username: username/email of the current user98        :param totp_secret: a unique shared secret of the user99        """100        tp = self._totp.from_source(totp_secret)101        return tp.to_uri(username)102 103    def get_totp_pretty_key(self, totp_secret: str) -> str:104        """Generate pretty key for manual input105 106        :param totp_secret: a unique shared secret of the user107 108        .. versionadded:: 4.0.0109        """110        tp = self._totp.from_source(totp_secret)111        return tp.pretty_key()112 113    def fetch_setup_values(self, totp: str, user: User) -> dict[str, str]:114        """Generate various values user needs to setup authenticator app.115            Returns dict with keys:116                'key': totp key117                'image': image as string (useful for <img src=xx>)118                'username: qrcode best practice119                'issuer': qrcode best practice120 121        .. versionadded:: 4.0.0122        """123 124        r = dict()125 126        # By convention, the URI should have the username that the user127        # logs in with.128        username = user.calc_username() or "Unknown"129        r["username"] = username130        r["key"] = self.get_totp_pretty_key(totp)131        r["issuer"] = self._totp.issuer132        r["image"] = self.generate_qrcode(username, totp)133        return r134 135    def generate_qrcode(self, username: str, totp: str) -> str:136        """Generate QRcode137         Using username, totp, generate the actual QRcode image.138         This method can be overridden to fine-tune how the image is created -139         such as size, color etc.140 141         It must return a string suitable for use in an <img src=xx> tag.142 143        .. versionadded:: 4.0.0144        """145        try:146            import qrcode147            import qrcode.image.svg148 149            image = qrcode.make(150                self.get_totp_uri(username, totp),151                image_factory=qrcode.image.svg.SvgImage,152            )153            with io.BytesIO() as virtual_file:154                image.save(virtual_file)155                image_as_str = base64.b64encode(virtual_file.getvalue()).decode("ascii")156 157            return f"data:image/svg+xml;base64,{image_as_str}"158        except ImportError:  # pragma: no cover159            # This should have been checked at app init.160            raise161 162    def generate_recovery_codes(self, number: int) -> list[str]:163        """Generate a set of secure passwords - used for 2FA recovery codes.164            # this is nice for english - but not for others165            return genphrase(entropy="fair", wordset="eff_short", sep="-",166             returns=number)167 168        .. versionadded:: 5.0.0169        """170        pwds = genword(length=12, charset="hex", returns=number)171        # make this a bit easier to type - 3 sets of 4 characters172        spwds = []173        for pwd in pwds:174            spwds.append(175                "-".join([pwd[i : i + 4] for i in range(0, len(pwd), 4)])  # noqa: E203176            )177        return spwds178 179    def get_last_counter(self, user: User) -> TotpMatch | None:180        """Implement this to fetch stored last_counter from cache.181 182        :param user: User model183        :return: last_counter as stored in set_last_counter()184        """185        return None186 187    def set_last_counter(self, user: User, tmatch: TotpMatch) -> None:188        """Implement this to cache last_counter.189 190        :param user: User model191        :param tmatch: a TotpMatch as returned from totp.verify()192        """193        pass194