From 8fc88b82530c29b4c1196d763fd210186451ac67 Mon Sep 17 00:00:00 2001
From: Son NK <>
Date: Thu, 9 Apr 2020 22:22:26 +0200
Subject: [PATCH] Set referral when creating User

---
 app/auth/views/facebook.py    |  7 +++++--
 app/auth/views/github.py      |  7 +++++--
 app/auth/views/google.py      |  7 +++++--
 app/auth/views/login_utils.py | 20 +++++++++++++++++++-
 app/auth/views/register.py    | 10 ++++++++--
 5 files changed, 42 insertions(+), 9 deletions(-)

diff --git a/app/auth/views/facebook.py b/app/auth/views/facebook.py
index 58756e36..cc15afe2 100644
--- a/app/auth/views/facebook.py
+++ b/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()
 
diff --git a/app/auth/views/github.py b/app/auth/views/github.py
index 73dc25b9..83a5af6e 100644
--- a/app/auth/views/github.py
+++ b/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)
diff --git a/app/auth/views/google.py b/app/auth/views/google.py
index eb92a8c6..495ec937 100644
--- a/app/auth/views/google.py
+++ b/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()
 
diff --git a/app/auth/views/login_utils.py b/app/auth/views/login_utils.py
index 246e22ca..d24a121b 100644
--- a/app/auth/views/login_utils.py
+++ b/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
diff --git a/app/auth/views/register.py b/app/auth/views/register.py
index aaf53b6b..ca692715 100644
--- a/app/auth/views/register.py
+++ b/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: