Refactor handling if cookies are disabled

app.py

@@ -31,12 +31,10 @@ def create_app(test_config=None):
 			user_groups = session['user_groups']
 		except (KeyError, OverflowError):
 			session.clear()
-			session['cookies_enabled'] = True
 			abort(401)
 		if datetime.datetime.now() - timestamp > datetime.timedelta(days=2) or \
 					client_id != request.headers['X-CLIENT-ID']:
 			session.clear()
-			session['cookies_enabled'] = True
 			abort(401)
 		resp = Response('Ok', 200)
 		resp.headers['OAUTH-USER-ID'] = user_id
@@ -51,14 +49,18 @@ def create_app(test_config=None):
 		return OAuth2Session(request.headers['X-CLIENT-ID'],
 			redirect_uri=request.headers['X-REDIRECT-URI'], **kwargs)
 
-	@app.route('/login')
-	def login():
-		# The cookies_enabled check prevents redirect loops:
-		# login (sets state) -> idp_authorize -> callback (no state set) -> login
+	@app.route('/cookiecheck')
+	def login_cookiecheck():
+		print(session)
 		if not session.get('cookies_enabled'):
 			session.clear()
 			session['cookies_enabled'] = True
 			abort(400, description='Enable cookies and reload two times to continue')
+		session.clear()
+		return redirect(url_for('login', url=request.values.get('url', '/')))
+
+	@app.route('/login')
+	def login():
 		client = get_oauth()
 		url, state = client.authorization_url(app.config['OAUTH2_AUTH_URL'])
 		session.clear()
@@ -76,7 +78,8 @@ def create_app(test_config=None):
 		if 'state' not in session:
 			session.clear()
 			session['cookies_enabled'] = True
-			return redirect(url_for('login', url=redirect_url))
+			# Redirect to login_cookiecheck to prevent redirect loop when cookies are disabled
+			return redirect(url_for('login_cookiecheck', url=redirect_url))
 		state = session['state']
 
 		client = get_oauth(state=state)

test_app.py

@@ -78,12 +78,8 @@ class TestCases(unittest.TestCase):
 	def test_auth_no_session(self):
 		r = self.client.get(path='/auth', headers=headers)
 		self.assertEqual(r.status_code, 401)
-		with self.client.session_transaction() as session:
-			self.assertEqual(session['cookies_enabled'], True)
 
 	def test_login(self):
-		with self.client.session_transaction() as session:
-			session['cookies_enabled'] = True
 		r = self.client.get(path='/login', query_string={'url': 'https://127.0.0.123:7654/app'}, headers=headers, follow_redirects=False)
 		self.assertEqual(r.status_code, 302)
 		url = urllib.parse.urlparse(r.location)
@@ -99,10 +95,6 @@ class TestCases(unittest.TestCase):
 			self.assertEqual(session['state'], qs['state'][0])
 			self.assertEqual(session['url'], 'https://127.0.0.123:7654/app')
 
-	def test_login_no_cookies(self):
-		r = self.client.get(path='/login', query_string={'url': 'https://127.0.0.123:7654/app'}, headers=headers, follow_redirects=False)
-		self.assertEqual(r.status_code, 400)
-
 	def test_callback(self):
 		code = 'testcode'
 		state = 'teststate'
@@ -130,7 +122,21 @@ class TestCases(unittest.TestCase):
 		r = self.client.get(path='/callback', headers=headers, query_string={'code': code, 'state': state}, follow_redirects=False)
 		self.assertEqual(r.status_code, 302)
 		url = urllib.parse.urlparse(r.location)
+		self.assertEqual(url.path, '/cookiecheck')
+		with self.client.session_transaction() as session:
+			self.assertEqual(session['cookies_enabled'], True)
+
+	def test_cookiecheck(self):
+		with self.client.session_transaction() as session:
+			session['cookies_enabled'] = True
+		r = self.client.get(path='/cookiecheck', headers=headers, follow_redirects=False)
+		self.assertEqual(r.status_code, 302)
+		url = urllib.parse.urlparse(r.location)
 		self.assertEqual(url.path, '/login')
+
+	def test_cookiecheck_no_session(self):
+		r = self.client.get(path='/cookiecheck', headers=headers, follow_redirects=False)
+		self.assertEqual(r.status_code, 400)
 		with self.client.session_transaction() as session:
 			self.assertEqual(session['cookies_enabled'], True)