Add CSRF protection using state to python-flask

python-flask/main.py

@@ -1,7 +1,8 @@
-from flask import Flask, render_template, request, session, redirect
+from flask import Flask, render_template, request, session, redirect, make_response
 import os, sys
 import requests
 import jwt
+import uuid
 
 CLIENT_ID = 'valtech.idp.testclient.local'
 CLIENT_SECRET = os.environ.get('CLIENT_SECRET')
@@ -27,12 +28,27 @@ def index():
 @app.route('/sign-in')
 def sign_in():
   if session.get('signed_in') != None: return redirect('/')
-  authorize_url = 'https://stage-id.valtech.com/oauth2/authorize?response_type=%s&client_id=%s&scope=%s' % ('code', CLIENT_ID, 'email openid')
-  return redirect(authorize_url)
+
+  # state is used for CSRF protection. the client generates a value and stores it
+  # for the user somewhere (in a cookie or in a session). it then passes the same value
+  # in the state parameter in the authorize request. IDP will mirror the state value
+  # to the redirect URI. the client should then make sure the state value it has stored
+  # matches what it receives in the callback
+  state = str(uuid.uuid4())
+
+  authorize_url = 'https://stage-id.valtech.com/oauth2/authorize?response_type=%s&client_id=%s&scope=%s&state=%s' % ('code', CLIENT_ID, 'email openid', state)
+
+  resp = make_response(redirect(authorize_url))
+  resp.set_cookie('python-flask-csrf', state)
+  return resp
 
 @app.route('/sign-in/callback')
 def sign_in_callback():
   code = request.args.get('code')
+  state = request.args.get('state')
+
+  if state != request.cookies.get('python-flask-csrf'):
+    raise Exception("Possible CSRF detected (state does not match stored state)")
 
   # as both scope openid and email was requested on authorize request above, the client
   # will receive both an access_token (according to OAuth 2) AND an id_token (according to OpenID Connect)
@@ -49,8 +65,10 @@ def sign_in_callback():
 
   session['signed_in'] = True
   session['email'] = user_info['email']
-
-  return redirect('/')
+
+  resp = make_response(redirect('/'))
+  resp.set_cookie('python-flask-csrf', '', expires=0)
+  return resp
 
 @app.route('/sign-out')
 def sign_out():