Adjustment to prevent exception in Google/Github authentication when local user cannot be created

app/views.py

@@ -17,7 +17,7 @@ from werkzeug import secure_filename
 from werkzeug.security import gen_salt
 
 from .models import User, Domain, Record, Server, History, Anonymous, Setting, DomainSetting
-from app import app, login_manager, github
+from app import app, login_manager, github, google
 from app.lib import utils
 
 
@@ -172,6 +172,13 @@ def register():
         return render_template('errors/404.html'), 404
 
 
+@app.route('/google/login')
+def google_login():
+    if not app.config.get('GOOGLE_OAUTH_ENABLE'):
+        return abort(400)
+    return google.authorize(callback=url_for('authorized', _external=True))
+
+
 @app.route('/github/login')
 def github_login():
     if not app.config.get('GITHUB_OAUTH_ENABLE'):
@@ -188,10 +195,34 @@ def login():
     BASIC_ENABLED = app.config['BASIC_ENABLED']
     SIGNUP_ENABLED = app.config['SIGNUP_ENABLED']
     GITHUB_ENABLE = app.config.get('GITHUB_OAUTH_ENABLE')
+    GOOGLE_ENABLE = app.config.get('GOOGLE_OAUTH_ENABLE')
 
     if g.user is not None and current_user.is_authenticated:
         return redirect(url_for('dashboard'))
 
+    if 'google_token' in session:
+        user_data = google.get('userinfo').data
+        first_name = user_data['given_name']
+        surname = user_data['family_name']
+        email = user_data['email']
+        user = User.query.filter_by(username=email).first()
+        if not user:
+            # create user
+            user = User(username=email,
+                        firstname=first_name,
+                        lastname=surname,
+                        plain_text_password=gen_salt(7),
+                        email=email)
+
+            result = user.create_local_user()
+            if not result['status']:
+                session.pop('google_token', None)
+                return redirect(url_for('login'))
+
+        session['user_id'] = user.id
+        login_user(user, remember = False)
+        return redirect(url_for('index'))
+
     if 'github_token' in session:
         me = github.get('user')
         user_info = me.data
@@ -201,7 +232,11 @@ def login():
             user = User(username=user_info['name'],
                         plain_text_password=gen_salt(7),
                         email=user_info['email'])
-            user.create_local_user()
+
+            result = user.create_local_user()
+            if not result['status']:
+                session.pop('github_token', None)
+                return redirect(url_for('login'))
 
         session['user_id'] = user.id
         login_user(user, remember = False)
@@ -210,6 +245,7 @@ def login():
     if request.method == 'GET':
         return render_template('login.html',
                                github_enabled=GITHUB_ENABLE,
+                               google_enabled=GOOGLE_ENABLE,
                                ldap_enabled=LDAP_ENABLED, login_title=LOGIN_TITLE,
                                basic_enabled=BASIC_ENABLED, signup_enabled=SIGNUP_ENABLED)
 
@@ -263,10 +299,10 @@ def login():
 
         try:
             result = user.create_local_user()
-            if result == True:
+            if result['status'] == True:
                 return render_template('login.html', username=username, password=password, ldap_enabled=LDAP_ENABLED, login_title=LOGIN_TITLE, basic_enabled=BASIC_ENABLED, signup_enabled=SIGNUP_ENABLED)
             else:
-                return render_template('register.html', error=result)
+                return render_template('register.html', error=result['msg'])
         except Exception as e:
             return render_template('register.html', error=e)
 
@@ -275,6 +311,7 @@ def login():
 def logout():
     session.pop('user_id', None)
     session.pop('github_token', None)
+    session.pop('google_token', None)
     logout_user()
     return redirect(url_for('login'))
 
@@ -628,14 +665,10 @@ def admin_createuser():
             return render_template('admin_createuser.html', user=user, blank_password=True)
 
         result = user.create_local_user();
+        if result['status']:
+            return redirect(url_for('admin_manageuser'))
 
-        if result == 'Email already existed':
-            return render_template('admin_createuser.html', user=user, duplicate_email=True)
-
-        if result == 'Username already existed':
-            return render_template('admin_createuser.html', user=user, duplicate_username=True)
-
-        return redirect(url_for('admin_manageuser'))
+        return render_template('admin_createuser.html', user=user, error=result['msg'])
 
 
 @app.route('/admin/manageuser', methods=['GET', 'POST'])