From c2bb6488e44376937d6cdb42a01d6ccf5a191b22 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Adri=C3=A0=20Casaj=C3=BAs?=
 <acasajus@users.noreply.github.com>
Date: Mon, 4 Jul 2022 16:09:36 +0200
Subject: [PATCH] Allow to login with proton to enter sudo mode (#1141)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* Allow to login with proton to enter sudo mode

* Updated wording

* lint

* Only enabled if the user has the account linked

* Add exit-sudo route for tests

Co-authored-by: Adrià Casajús <adria.casajus@proton.ch>
---
 app/auth/views/fido.py              |  2 ++
 app/auth/views/login_utils.py       |  2 ++
 app/dashboard/views/enter_sudo.py   | 14 +++++++++++++-
 app/internal/__init__.py            |  1 +
 app/internal/exit_sudo.py           | 10 ++++++++++
 templates/dashboard/enter_sudo.html | 13 +++++++++++++
 6 files changed, 41 insertions(+), 1 deletion(-)
 create mode 100644 app/internal/exit_sudo.py

diff --git a/app/auth/views/fido.py b/app/auth/views/fido.py
index c65a6ff5..445fd83f 100644
--- a/app/auth/views/fido.py
+++ b/app/auth/views/fido.py
@@ -1,5 +1,6 @@
 import json
 import secrets
+from time import time
 
 import webauthn
 from flask import (
@@ -107,6 +108,7 @@ def fido():
             Session.commit()
             del session[MFA_USER_ID]
 
+            session["sudo_time"] = int(time())
             login_user(user)
             flash(f"Welcome back!", "success")
 
diff --git a/app/auth/views/login_utils.py b/app/auth/views/login_utils.py
index 96bd0ab5..f6563e10 100644
--- a/app/auth/views/login_utils.py
+++ b/app/auth/views/login_utils.py
@@ -1,3 +1,4 @@
+from time import time
 from typing import Optional
 
 from flask import session, redirect, url_for, request
@@ -31,6 +32,7 @@ def after_login(user, next_url):
     else:
         LOG.d("log user %s in", user)
         login_user(user)
+        session["sudo_time"] = int(time())
 
         # User comes to login page from another page
         if next_url:
diff --git a/app/dashboard/views/enter_sudo.py b/app/dashboard/views/enter_sudo.py
index d45f5c00..6e937dcc 100644
--- a/app/dashboard/views/enter_sudo.py
+++ b/app/dashboard/views/enter_sudo.py
@@ -8,6 +8,8 @@ from wtforms import PasswordField, validators
 
 from app.dashboard.base import dashboard_bp
 from app.log import LOG
+from app.models import PartnerUser
+from app.proton.utils import is_connect_with_proton_enabled, get_proton_partner
 from app.utils import sanitize_next_url
 
 _SUDO_GAP = 900
@@ -39,8 +41,18 @@ def enter_sudo():
         else:
             flash("Incorrect password", "warning")
 
+    proton_enabled = is_connect_with_proton_enabled()
+    if proton_enabled:
+        # Only for users that have the account linked
+        partner_user = PartnerUser.get_by(user_id=current_user.id)
+        if not partner_user or partner_user.partner_id != get_proton_partner().id:
+            proton_enabled = False
+
     return render_template(
-        "dashboard/enter_sudo.html", password_check_form=password_check_form
+        "dashboard/enter_sudo.html",
+        password_check_form=password_check_form,
+        next=request.args.get("next"),
+        connect_with_proton=proton_enabled,
     )
 
 
diff --git a/app/internal/__init__.py b/app/internal/__init__.py
index c5ec6c12..c92d4670 100644
--- a/app/internal/__init__.py
+++ b/app/internal/__init__.py
@@ -1 +1,2 @@
 from .integrations import set_enable_proton_cookie
+from .exit_sudo import exit_sudo_mode
diff --git a/app/internal/exit_sudo.py b/app/internal/exit_sudo.py
new file mode 100644
index 00000000..cf10a155
--- /dev/null
+++ b/app/internal/exit_sudo.py
@@ -0,0 +1,10 @@
+from flask import session, redirect, url_for, flash
+
+from app.internal.base import internal_bp
+
+
+@internal_bp.route("/exit-sudo-mode")
+def exit_sudo_mode():
+    session["sudo_time"] = 0
+    flash("Exited sudo mode", "info")
+    return redirect(url_for("dashboard.index"))
diff --git a/templates/dashboard/enter_sudo.html b/templates/dashboard/enter_sudo.html
index 14df9da1..d176ca08 100644
--- a/templates/dashboard/enter_sudo.html
+++ b/templates/dashboard/enter_sudo.html
@@ -16,6 +16,19 @@
         {{ render_field_errors(password_check_form.password) }}
         <button class="btn btn-lg btn-danger mt-2">Submit</button>
       </form>
+      {% if connect_with_proton %}
+
+        <div class="my-3">
+          <p>
+            Alternatively you can use your Proton credentials to ensure it's you.
+          </p>
+        </div>
+        <a class="btn btn-primary btn-block mt-2 proton-button w-25"
+           href="{{ url_for("auth.proton_login", next=next) }}">
+          <img class="mr-2" src="/static/images/proton.svg" />
+          Authenticate with Proton
+        </a>
+      {% endif %}
     </div>
   </div>
 {% endblock %}