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 by sending OTP through email"""10 11from flask import url_for, session, Response, render_template, current_app, \12    flash13from flask_babel import gettext as _14from flask_login import current_user15from flask_security import send_mail16 17import config18from pgadmin.utils.csrf import pgCSRFProtect19from .registry import BaseMFAuth20from .utils import ValidationException, mfa_add, fetch_auth_option21from pgadmin.utils.constants import MessageType22 23 24def __generate_otp() -> str:25    """26    Generate a six-digits one-time-password (OTP) for the current user.27 28    Returns:29        str: A six-digits OTP for the current user30    """31    import time32    import codecs33    import secrets34 35    code = codecs.encode("{}{}{}".format(36        time.time(), current_user.username, secrets.choice(range(1000, 9999))37    ).encode(), "hex")38 39    res = 040    idx = 041 42    while idx < len(code):43        res += int((code[idx:idx + 6]).decode('utf-8'), base=16)44        res %= 100000045        idx += 546 47    return str(res).zfill(6)48 49 50def _send_code_to_email(_email: str = None) -> (bool, int, str):51    """52    Send the code to the email address, provided in the argument or to the53    email address of the current user, provided during the registration.54 55    Args:56        _email (str, optional): Email Address, where to send the OTP code.57                                Defaults to None.58 59    Returns:60        (bool, int, str): Returns a set as (failed?, HTTP Code, message string)61                          If 'failed?' is True, message contains the error62                          message for the user, else it contains the success63                          message for the user to consume.64    """65 66    if not current_user.is_authenticated:67        return False, 401, _("Not accessible")68 69    if _email is None:70        _email = getattr(current_user, 'email', None)71 72    if _email is None:73        return False, 401, _("No email address is available.")74 75    try:76        session["mfa_email_code"] = __generate_otp()77        subject = getattr(config, 'MFA_EMAIL_SUBJECT', None)78 79        if subject is None:80            subject = _("{} - Verification Code").format(config.APP_NAME)81 82        send_mail(83            subject,84            _email,85            "send_email_otp",86            user=current_user,87            code=session["mfa_email_code"]88        )89    except OSError as ose:90        current_app.logger.exception(ose)91        return False, 503, _("Failed to send the code to email.") + \92            "\n" + str(ose)93 94    message = _(95        "A verification code was sent to {}. Check your email and enter "96        "the code."97    ).format(_mask_email(_email))98 99    return True, 200, message100 101 102def _mask_email(_email: str) -> str:103    """104 105    Args:106        _email (str): Email address to be masked107 108    Returns:109        str: Masked email address110    """111    import re112    email_split = re.split('@', _email)113    username, domain = email_split114    domain_front, *domain_back_list = re.split('[.]', domain)115    users = re.split('[.]', username)116 117    def _mask_except_first_char(_str: str) -> str:118        """119        Mask all characters except first character of the input string.120        Args:121            _str (str): Input string to be masked122 123        Returns:124            str: Masked string125        """126        return _str[0] + '*' * (len(_str) - 1)127 128    return '.'.join([_mask_except_first_char(user) for user in users]) + \129        '@' + _mask_except_first_char(domain_front) + '.' + \130        '.'.join(domain_back_list)131 132 133def send_email_code() -> Response:134    """135    Send the code to the users' email address, stored during the registration.136 137    Raises:138        ValidationException: Raise this exception when user is not registered139        for this authentication method.140 141    Returns:142        Flask.Response: Response containing the HTML portion after sending the143                        code to the registered email address of the user.144    """145 146    options, found = fetch_auth_option(EMAIL_AUTH_METHOD)147 148    if found is False:149        raise ValidationException(_(150            "User has not registered for email authentication"151        ))152 153    success, http_code, message = _send_code_to_email(options)154 155    if success is False:156        return Response(message, http_code, mimetype='text/html')157 158    return dict(message=message)159 160 161@pgCSRFProtect.exempt162def javascript() -> Response:163    """164    Returns the javascript code for the email authentication method.165 166    Returns:167        Flask.Response: Response object conataining the javscript code for the168                        email auth method.169    """170    if not current_user.is_authenticated:171        return Response(_("Not accessible"), 401, mimetype="text/plain")172 173    return Response(render_template(174        "mfa/email.js", _=_, url_for=url_for,175    ), 200, mimetype="text/javascript")176 177 178EMAIL_AUTH_METHOD = 'email'179 180 181def email_authentication_label():182    return _('Email Authentication')183 184 185class EmailAuthentication(BaseMFAuth):186 187    @property188    def name(self):189        return EMAIL_AUTH_METHOD190 191    @property192    def label(self):193        return email_authentication_label()194 195    def validate(self, **kwargs):196        code = kwargs.get('code', None)197        email_otp = session.get("mfa_email_code", None)198        if code is not None and email_otp is not None and code == email_otp:199            session.pop("mfa_email_code")200            return201        raise ValidationException("Invalid code")202 203    def validation_view(self):204        session.pop("mfa_email_code", None)205        return dict(206            description=_("Verify with Email Authentication"),207            button_label=_("Send Code"),208            button_label_sending=_("Sending Code...")209        )210 211    def _registration_view(self):212        email = getattr(current_user, 'email', '')213        return dict(214            label=email_authentication_label(),215            auth_method=EMAIL_AUTH_METHOD,216            description=_("Enter the email address to send a code"),217            email_address_placeholder=_("Email address"),218            email_address=email,219            note_label=_("Note"),220            note=_(221                "This email address will only be used for two factor "222                "authentication purposes. The email address for the user "223                "account will not be changed."224            ),225        )226 227    def _registration_view_after_code_sent(self, _form_data):228 229        session['mfa_email_id'] = _form_data.get('send_to', None)230        success, http_code, message = _send_code_to_email(231            session['mfa_email_id']232        )233 234        if success is False:235            flash(message, MessageType.ERROR)236            return None237 238        return dict(239            label=email_authentication_label(),240            auth_method=EMAIL_AUTH_METHOD,241            message=message,242            otp_placeholder=_("Enter code here"),243            http_code=http_code,244        )245 246    def registration_view(self, _form_data):247 248        if 'validate' in _form_data:249            if _form_data['validate'] == 'send_code':250                return self._registration_view_after_code_sent(_form_data)251 252            code = _form_data.get('code', 'unknown')253 254            if code is not None and \255                    code == session.get("mfa_email_code", None) and \256                    session.get("mfa_email_id", None) is not None:257                mfa_add(EMAIL_AUTH_METHOD, session['mfa_email_id'])258 259                flash(_(260                    "Email Authentication registered successfully."261                ), MessageType.SUCCESS)262 263                session.pop('mfa_email_code', None)264 265                return None266 267            flash(_('Invalid code'), MessageType.ERROR)268 269        return self._registration_view()270 271    def register_url_endpoints(self, blueprint):272        blueprint.add_url_rule(273            "/send_email_code", "send_email_code", send_email_code,274            methods=("POST", )275        )276        blueprint.add_url_rule(277            "/email.js", "email_js", javascript, methods=("GET", )278        )279 280    @property281    def icon(self):282        return url_for("mfa.static", filename="images/email_lock.svg")283 284    @property285    def validate_script(self):286        return url_for("mfa.email_js")287