codekingpro/portable-devtools

1##############################################################################2#3# pgAdmin 4 - PostgreSQL Tools4#5# Copyright (C) 2013 - 2024, The pgAdmin Development Team6# This software is released under the PostgreSQL Licence7#8##############################################################################9"""Multi-factor Authentication implementation for Time-based One-Time Password10(TOTP) applications"""11 12import base6413from io import BytesIO14from typing import Union15 16from flask import url_for, session, flash17from flask_babel import gettext as _18from flask_login import current_user19import pyotp20import qrcode21 22import config23from pgadmin.model import UserMFA24 25from .registry import BaseMFAuth26from .utils import ValidationException, fetch_auth_option, mfa_add27from pgadmin.utils.constants import MessageType28 29 30_TOTP_AUTH_METHOD = "authenticator"31_TOTP_AUTHENTICATOR = _("Authenticator App")32_OTP_PLACEHOLDER = _("Enter code")33 34 35class TOTPAuthenticator(BaseMFAuth):36    """37    Authenction class for TOTP based authentication.38 39    Base Class: BaseMFAuth40    """41 42    @classmethod43    def __create_topt_for_currentuser(cls) -> pyotp.TOTP:44        """45        Create the TOPT object using the secret stored for the current user in46        the configuration database.47 48        Assumption: Configuration database is not modified by anybody manually,49                    and removed the secrete for the current user.50 51        Raises:52            ValidationException: Raises when user is not registered for this53                                 authenction method.54 55        Returns:56            pyotp.TOTP: TOTP object for the current user (if registered)57        """58        options, found = fetch_auth_option(_TOTP_AUTH_METHOD)59 60        if found is False:61            raise ValidationException(_(62                "User has not registered the Time-based One-Time Password "63                "(TOTP) Authenticator for authentication."64            ))65 66        if options is None or options == '':67            raise ValidationException(_(68                "User does not have valid HASH to generate the OTP."69            ))70 71        return pyotp.TOTP(options)72 73    @property74    def name(self) -> str:75        """76        Name of the authetication method for internal presentation.77 78        Returns:79            str: Short name for this authentication method80        """81        return _TOTP_AUTH_METHOD82 83    @property84    def label(self) -> str:85        """86        Label for the UI for this authentication method.87 88        Returns:89            str: User presentable string for this auth method90        """91        return _(_TOTP_AUTHENTICATOR)92 93    @property94    def icon(self) -> str:95        """96        Property for the icon url string for this auth method, to be used on97        the authentication or registration page.98 99        Returns:100            str: url for the icon representation for this auth method101        """102        return url_for("mfa.static", filename="images/totp_lock.svg")103 104    def validate(self, **kwargs):105        """106        Validate the code sent using the HTTP request.107 108        Raises:109            ValidationException: Raises when code is not valid110        """111        code = kwargs.get('code', None)112        totp = TOTPAuthenticator.__create_topt_for_currentuser()113 114        if totp.verify(code) is False:115            raise ValidationException("Invalid Code")116 117    def validation_view(self) -> dict:118        """119        Generate the portion of the view to render on the authentication page120 121        Returns:122            str: Authentication view as a string123        """124        return dict(125            auth_description=_(126                "Enter the code shown in your authenticator application for "127                "TOTP (Time-based One-Time Password)"128            ),129            otp_placeholder=_OTP_PLACEHOLDER,130        )131 132    def _registration_view(self) -> dict:133        """134        Internal function to generate a view for the registration page.135 136        View will contain the QRCode image for the TOTP based authenticator137        applications to scan.138 139        Returns:140            str: Registration view with QRcode for TOTP based applications141        """142 143        option = session.pop('mfa_authenticator_opt', None)144        if option is None:145            option = pyotp.random_base32()146        session['mfa_authenticator_opt'] = option147        totp = pyotp.TOTP(option)148 149        uri = totp.provisioning_uri(150            current_user.username, issuer_name=getattr(151                config, "APP_NAME", "pgAdmin 4"152            )153        )154 155        img = qrcode.make(uri)156        buffered = BytesIO()157        img.save(buffered)158        img_base64 = base64.b64encode(buffered.getvalue())159 160        return dict(161            auth_title=_(_TOTP_AUTHENTICATOR),162            auth_method=_TOTP_AUTH_METHOD,163            image=img_base64.decode("utf-8"),164            qrcode_alt_text=_("TOTP Authenticator QRCode"),165            auth_description=_(166                "Scan the QR code and the enter the code from the "167                "TOTP Authenticator application"168            ), otp_placeholder=_OTP_PLACEHOLDER169        )170 171    def registration_view(self, form_data) -> Union[str, None]:172        """173        Returns the registration view for this authentication method.174 175        It is also responsible for validating the code during the registration.176 177        Args:178            form_data (dict): Form data as a dictionary sent from the179                              registration page for rendering or validation of180                              the code.181 182        Returns:183            str: Registration view for the 'authenticator' method if it is not184                 a request for the validation of the code or the code sent is185                 not a valid TOTP code, otherwise - it will return None.186        """187 188        if 'VALIDATE' not in form_data:189            return self._registration_view()190 191        code = form_data.get('code', None)192        authenticator_opt = session.get('mfa_authenticator_opt', None)193        if authenticator_opt is None or \194                pyotp.TOTP(authenticator_opt).verify(code) is False:195            flash(_("Failed to validate the code"), MessageType.ERROR)196            return self._registration_view()197 198        mfa_add(_TOTP_AUTH_METHOD, authenticator_opt)199        flash(_(200            "TOTP Authenticator registered successfully for authentication."201        ), MessageType.SUCCESS)202        session.pop('mfa_authenticator_opt', None)203 204        return None205