Hide proton integration behind cookie (#1092)

* Hide proton integration behind cookie

* Make cookie name configurable via config

app/auth/views/login.py

@@ -5,11 +5,11 @@ from wtforms import StringField, validators
 
 from app.auth.base import auth_bp
 from app.auth.views.login_utils import after_login
-from app.config import CONNECT_WITH_PROTON
 from app.events.auth_event import LoginEvent
 from app.extensions import limiter
 from app.log import LOG
 from app.models import User
+from app.proton.utils import is_connect_with_proton_enabled
 from app.utils import sanitize_email, sanitize_next_url
 
 
@@ -68,5 +68,5 @@ def login():
         form=form,
         next_url=next_url,
         show_resend_activation=show_resend_activation,
-        connect_with_proton=CONNECT_WITH_PROTON,
+        connect_with_proton=is_connect_with_proton_enabled(),
     )

app/config.py

@@ -246,6 +246,7 @@ PROTON_VALIDATE_CERTS = "PROTON_VALIDATE_CERTS" in os.environ
 CONNECT_WITH_PROTON = "CONNECT_WITH_PROTON" in os.environ
 PROTON_EXTRA_HEADER_NAME = os.environ.get("PROTON_EXTRA_HEADER_NAME")
 PROTON_EXTRA_HEADER_VALUE = os.environ.get("PROTON_EXTRA_HEADER_VALUE")
+CONNECT_WITH_PROTON_COOKIE_NAME = os.environ.get("CONNECT_WITH_PROTON_COOKIE_NAME")
 
 # in seconds
 AVATAR_URL_EXPIRATION = 3600 * 24 * 7  # 1h*24h/d*7d=1week

app/dashboard/views/setting.py

@@ -21,7 +21,6 @@ from app.config import (
     URL,
     FIRST_ALIAS_DOMAIN,
     ALIAS_RANDOM_SUFFIX_LENGTH,
-    CONNECT_WITH_PROTON,
 )
 from app.dashboard.base import dashboard_bp
 from app.db import Session
@@ -50,6 +49,7 @@ from app.models import (
     AppleSubscription,
     PartnerUser,
 )
+from app.proton.utils import is_connect_with_proton_enabled
 from app.proton.proton_callback_handler import get_proton_partner
 from app.utils import random_string, sanitize_email
 
@@ -376,7 +376,7 @@ def setting():
         coinbase_sub=coinbase_sub,
         FIRST_ALIAS_DOMAIN=FIRST_ALIAS_DOMAIN,
         ALIAS_RAND_SUFFIX_LENGTH=ALIAS_RANDOM_SUFFIX_LENGTH,
-        connect_with_proton=CONNECT_WITH_PROTON,
+        connect_with_proton=is_connect_with_proton_enabled(),
         proton_linked_account=proton_linked_account,
     )
 

app/internal/__init__.py

@@ -0,0 +1 @@
+from .integrations import set_enable_proton_cookie

app/internal/base.py

@@ -0,0 +1,8 @@
+from flask import Blueprint
+
+internal_bp = Blueprint(
+    name="internal",
+    import_name=__name__,
+    url_prefix="/internal",
+    template_folder="templates",
+)

app/internal/integrations.py

@@ -0,0 +1,25 @@
+import arrow
+from app.config import CONNECT_WITH_PROTON_COOKIE_NAME, URL
+from flask import make_response, redirect, url_for
+from flask_login import current_user
+from .base import internal_bp
+
+
+@internal_bp.route("/integrations/proton")
+def set_enable_proton_cookie():
+    if current_user.is_authenticated:
+        redirect_url = url_for("dashboard.index")
+    else:
+        redirect_url = url_for("auth.login")
+
+    response = make_response(redirect(redirect_url))
+    if CONNECT_WITH_PROTON_COOKIE_NAME:
+        response.set_cookie(
+            CONNECT_WITH_PROTON_COOKIE_NAME,
+            value="true",
+            expires=arrow.now().shift(days=30).datetime,
+            secure=True if URL.startswith("https") else False,
+            httponly=True,
+            samesite="Lax",
+        )
+    return response

app/proton/utils.py

@@ -0,0 +1,12 @@
+from app.config import CONNECT_WITH_PROTON, CONNECT_WITH_PROTON_COOKIE_NAME
+from flask import request
+
+
+def is_connect_with_proton_enabled() -> bool:
+    if CONNECT_WITH_PROTON:
+        return True
+    if CONNECT_WITH_PROTON_COOKIE_NAME and request.cookies.get(
+        CONNECT_WITH_PROTON_COOKIE_NAME
+    ):
+        return True
+    return False

example.env

@@ -114,6 +114,7 @@ WORDS_FILE_PATH=local_data/test_words.txt
 # PROTON_BASE_URL=to_fill
 # PROTON_VALIDATE_CERTS=true
 # CONNECT_WITH_PROTON=true
+# CONNECT_WITH_PROTON_COOKIE_NAME=to_fill
 
 # Flask profiler
 # FLASK_PROFILER_PATH=/tmp/flask-profiler.sql

server.py

@@ -77,6 +77,7 @@ from app.discover.base import discover_bp
 from app.email_utils import send_email, render
 from app.extensions import login_manager, limiter
 from app.fake_data import fake_data
+from app.internal.base import internal_bp
 from app.jose_utils import get_jwk_key
 from app.log import LOG
 from app.models import (
@@ -221,6 +222,7 @@ def register_blueprints(app: Flask):
     app.register_blueprint(onboarding_bp)
 
     app.register_blueprint(discover_bp)
+    app.register_blueprint(internal_bp)
     app.register_blueprint(api_bp)