Modified the python server to support SAML account merge/upgrade upon first SAML-type login of an existing user.

2 files changed, 67 insertions, 24 deletions

diff --git a/server/handlers.py b/server/handlers.py
index 30f52f1..da257c1 100644
--- a/server/handlers.py
+++ b/server/handlers.py
@@ -303,14 +303,20 @@ class SamlLogin(CodeqService):
         js = request.data
         saml_data = js.get('saml_data')
         gui_lang = js.get('gui_lang', 'en')
+        upgrade_account = js.get('upgrade_account')
+        upgrade_password = js.get('upgrade_password')
         if saml_data is None:
-            request.reply({'code': 1, 'message': 'SAML user data not specified'})
+            request.reply({'code': 3, 'message': 'SAML user data not specified'})
         else:
             session = request.session
             try:
-                name, email, date_joined, last_login = session.saml_login_or_signup(saml_data, gui_lang)
-            except Exception as e:
-                request.reply({'code': 2, 'message': 'SAML login failed: ' + str(e)})
+                name, email, date_joined, last_login = session.saml_login_or_signup(saml_data, gui_lang, upgrade_account, upgrade_password)
+            except AccountMergeRequired as amr:
+                request.reply({'code': 1, 'message': 'Retry with "upgrade_account" and "upgrade_password"', 'username': amr.username})
+            except AuthenticationFailed as af:
+                request.reply({'code': 4, 'message': 'Password is incorrect'})
+#            except Exception as e:
+#                request.reply({'code': 2, 'message': 'SAML login failed: ' + str(e)})
             else:
                 request.reply({'code': 0, 'message': 'OK', 'name': name, 'email': email, 'joined': date_joined.isoformat(), 'last-login': last_login.isoformat(), 'settings': session.get_settings()})
 
diff --git a/server/user_session.py b/server/user_session.py
index 02e1bd4..8a2e88e 100644
--- a/server/user_session.py
+++ b/server/user_session.py
@@ -14,7 +14,7 @@ import uuid
 
 import db
 from db.models import Problem, Solution
-from errors.session import NoSuchSession, AuthenticationFailed, PasswordChangeFailed, UserExists, SignupFailed, NotLoggedIn
+from errors.session import *
 import server
 from server.problems import load_language, load_group
 import psycopg2.extras
@@ -117,41 +117,78 @@ class UserSession(object):
             finally:
                 db.return_connection(conn)
 
-    def saml_login_or_signup(self, saml_data, gui_lang):
-        uuid = saml_data.get('schacUUID')
-        if uuid is None:
-            raise AuthenticationFailed('SAML data does not contain schacUUID')
+    def saml_login_or_signup(self, saml_data, gui_lang, can_upgrade_account, upgrade_password):
+        #uuid = saml_data.get('schacUUID')
+        #if uuid is None:
+        #    raise AuthenticationFailed('SAML data does not contain schacUUID')
         name = saml_data.get('displayName')
-        email = saml_data.get('mail')
+        #email = saml_data.get('mail')
+        eduPersonPrincipalName = saml_data.get('eduPersonPrincipalName')
+        if eduPersonPrincipalName is None:
+            raise AuthenticationFailed('SAML data does not contain eduPersonPrincipalName')
         with self._access_lock:
             now = datetime.datetime.utcnow()
             conn = db.get_connection()
             try:
                 cur = conn.cursor()
                 try:
-                    cur.execute('update codeq_user set name = %s, email = %s, saml_data = %s, last_login = %s where username = %s and saml_data is not null returning id, gui_lang, date_joined, robot_address, gui_layout', (name, email, psycopg2.extras.Json(saml_data), str(now), uuid))
-                    row = cur.fetchone()
-                    if row:
-                        self.uid = row[0]
-                        self.username = uuid
-                        self.settings = {'gui_lang': row[1], 'robot_address': row[3], 'gui_layout': row[4]}
-                        return name, email, row[2], now
+                    cur.execute('select id, gui_lang, date_joined, robot_address, gui_layout, username, password from codeq_user where email = %s for update', (eduPersonPrincipalName,))
+                    data = cur.fetchone()  # go through all the records and prefer the one without a password, that one will be the SAML account
+                    if data:
+                        stored_password = data[6]
+                        row = cur.fetchone()
+                        while row:
+                            if row[6] is None:
+                                data = row
+                            row = cur.fetchone()
+                    else:
+                        stored_password = None
+                    if data and ((not stored_password) or (can_upgrade_account is None) or can_upgrade_account):
+                        # the account already exists
+                        if stored_password:
+                            # the account has a password: an upgrade must be made
+                            if can_upgrade_account is None:
+                                raise AccountMergeRequired(data[5])
+                            # can_upgrade_account can only be True here
+                            if not verify_password(upgrade_password, stored_password):
+                                raise AuthenticationFailed('Password mismatch')
+                        cur.execute('update codeq_user set password = null, name = %s, email = %s, saml_data = %s, last_login = %s where id = %s', (name, eduPersonPrincipalName, psycopg2.extras.Json(saml_data), str(now), data[0]))
+                        self.uid = data[0]
+                        self.username = data[5]
+                        self.settings = {'gui_lang': gui_lang, 'robot_address': data[3], 'gui_layout': data[4]}
+                        date_joined = data[2]
                     else:
-                        cur.execute('insert into codeq_user (username, name, email, is_admin, is_active, date_joined, last_login, gui_lang, saml_data) values (%s, %s, %s, %s, %s, %s, %s, %s, %s) returning id', (uuid, name, email, False, True, str(now), str(now), gui_lang, psycopg2.extras.Json(saml_data)))
+                        # a new account required
+                        cur.execute('insert into codeq_user (username, name, email, is_admin, is_active, date_joined, last_login, gui_lang, saml_data) values (%s, %s, %s, %s, %s, %s, %s, %s, %s) returning id', (eduPersonPrincipalName, name, eduPersonPrincipalName, False, True, str(now), str(now), gui_lang, psycopg2.extras.Json(saml_data)))
                         row = cur.fetchone()
                         if row is None:
                             raise SignupFailed('Sign-up failed')
                         self.uid = row[0]
-                        self.username = uuid
+                        self.username = eduPersonPrincipalName
                         self.settings = {'gui_lang': gui_lang, 'robot_address': None, 'gui_layout': None}
-                        return name, email, now, now
+                        date_joined = now
+
+                    # cur.execute('update codeq_user set name = %s, email = %s, saml_data = %s, last_login = %s where username = %s and saml_data is not null returning id, gui_lang, date_joined, robot_address, gui_layout', (name, email, psycopg2.extras.Json(saml_data), str(now), uuid))
+                    # row = cur.fetchone()
+                    # if row:
+                    #     self.uid = row[0]
+                    #     self.username = uuid
+                    #     self.settings = {'gui_lang': row[1], 'robot_address': row[3], 'gui_layout': row[4]}
+                    #     return name, email, row[2], now
+                    # else:
+                    #     cur.execute('insert into codeq_user (username, name, email, is_admin, is_active, date_joined, last_login, gui_lang, saml_data) values (%s, %s, %s, %s, %s, %s, %s, %s, %s) returning id', (uuid, name, email, False, True, str(now), str(now), gui_lang, psycopg2.extras.Json(saml_data)))
+                    #     row = cur.fetchone()
+                    #     if row is None:
+                    #         raise SignupFailed('Sign-up failed')
+                    #     self.uid = row[0]
+                    #     self.username = uuid
+                    #     self.settings = {'gui_lang': gui_lang, 'robot_address': None, 'gui_layout': None}
+                    #     return name, email, now, now
                 finally:
                     cur.close()
+                conn.commit()
+                return name, eduPersonPrincipalName, date_joined, now
             finally:
-                try:
-                    conn.commit()
-                except:
-                    pass
                 db.return_connection(conn)
 
     def logout(self):