Set referral when creating User

app/auth/views/facebook.py

@@ -15,7 +15,7 @@ from app.config import (
 from app.extensions import db
 from app.log import LOG
 from app.models import User, SocialAuth
-from .login_utils import after_login
+from .login_utils import after_login, get_referral
 from ...email_utils import can_be_used_as_personal_email, email_already_used
 
 _authorization_base_url = "https://www.facebook.com/dialog/oauth"
@@ -118,7 +118,10 @@ def facebook_callback():
 
         LOG.d("create facebook user with %s", facebook_user_data)
         user = User.create(
-            email=email.lower(), name=facebook_user_data["name"], activated=True
+            email=email.lower(),
+            name=facebook_user_data["name"],
+            activated=True,
+            referral=get_referral(),
         )
         db.session.flush()
 

app/auth/views/github.py

@@ -4,7 +4,7 @@ from requests_oauthlib import OAuth2Session
 
 from app import email_utils
 from app.auth.base import auth_bp
-from app.auth.views.login_utils import after_login
+from app.auth.views.login_utils import after_login, get_referral
 from app.config import GITHUB_CLIENT_ID, GITHUB_CLIENT_SECRET, URL, DISABLE_REGISTRATION
 from app.email_utils import can_be_used_as_personal_email, email_already_used
 from app.extensions import db
@@ -100,7 +100,10 @@ def github_callback():
 
         LOG.d("create github user")
         user = User.create(
-            email=email.lower(), name=github_user_data.get("name") or "", activated=True
+            email=email.lower(),
+            name=github_user_data.get("name") or "",
+            activated=True,
+            referral=get_referral(),
         )
         db.session.commit()
         login_user(user)

app/auth/views/google.py

@@ -9,7 +9,7 @@ from app.extensions import db
 from app.log import LOG
 from app.models import User, File, SocialAuth
 from app.utils import random_string
-from .login_utils import after_login
+from .login_utils import after_login, get_referral
 from ...email_utils import can_be_used_as_personal_email, email_already_used
 
 _authorization_base_url = "https://accounts.google.com/o/oauth2/v2/auth"
@@ -103,7 +103,10 @@ def google_callback():
 
         LOG.d("create google user with %s", google_user_data)
         user = User.create(
-            email=email.lower(), name=google_user_data["name"], activated=True
+            email=email.lower(),
+            name=google_user_data["name"],
+            activated=True,
+            referral=get_referral(),
         )
         db.session.flush()
 

app/auth/views/login_utils.py

@@ -1,8 +1,11 @@
-from flask import session, redirect, url_for
+from typing import Optional
+
+from flask import session, redirect, url_for, request
 from flask_login import login_user
 
 from app.config import MFA_USER_ID
 from app.log import LOG
+from app.models import Referral
 
 
 def after_login(user, next_url):
@@ -28,3 +31,18 @@ def after_login(user, next_url):
         else:
             LOG.debug("redirect user to dashboard")
             return redirect(url_for("dashboard.index"))
+
+
+# name of the cookie that stores the referral code
+_REFERRAL_COOKIE = "slref"
+
+
+def get_referral() -> Optional[Referral]:
+    """Get the eventual referral stored in cookie"""
+    # whether user arrives via a referral
+    referral = None
+    if request.cookies:
+        ref_code = request.cookies.get(_REFERRAL_COOKIE)
+        referral = Referral.get_by(code=ref_code)
+
+    return referral

app/auth/views/register.py

@@ -5,11 +5,12 @@ from wtforms import StringField, validators
 
 from app import email_utils, config
 from app.auth.base import auth_bp
+from app.auth.views.login_utils import get_referral
 from app.config import URL, DISABLE_REGISTRATION
 from app.email_utils import can_be_used_as_personal_email, email_already_used
 from app.extensions import db
 from app.log import LOG
-from app.models import User, ActivationCode
+from app.models import User, ActivationCode, Referral
 from app.utils import random_string, encode_url
 
 
@@ -43,7 +44,12 @@ def register():
                 flash(f"Email {email} already used", "error")
             else:
                 LOG.debug("create user %s", form.email.data)
-                user = User.create(email=email, name="", password=form.password.data)
+                user = User.create(
+                    email=email,
+                    name="",
+                    password=form.password.data,
+                    referral=get_referral(),
+                )
                 db.session.commit()
 
                 try: