csrf demo

Author: greyshell

web/dblib/mysql_db.py

@@ -37,5 +37,6 @@ def conn(self):
 
     @staticmethod
     def close():
-        MySQLdb.instance.conn.close()
-        MySQLdb.instance = None
+        if MySQLdb.instance is not None:
+            MySQLdb.instance.conn.close()
+            MySQLdb.instance = None

web/login/app.py

@@ -0,0 +1,87 @@
+#!/usr/bin/env python3
+
+# author: greyshell
+
+import os
+
+from flask import Flask, render_template, session, redirect, url_for, jsonify
+from flask import request
+from flask_wtf.csrf import CSRFProtect
+from web.login.models import *
+
+templates_path = os.path.abspath(
+        './template/'
+)
+app = Flask(__name__, template_folder=templates_path)
+app.config.update(dict(
+        SECRET_KEY="woopie"
+))
+csrf = CSRFProtect(app)
+
+
+def get_logged_in_user():
+    if "userid" not in session:
+        return None
+    return User.logged_user(session['userid'])
+
+
+@app.before_request
+def before_request_func():
+    MySQLdb()  # create database
+
+
+@app.after_request
+def after_request_func(response):
+    MySQLdb.close()
+    return response
+
+
+@app.route('/', methods=['GET', 'POST'])
+@csrf.exempt
+def index():
+    user = get_logged_in_user()
+    if user is not None:
+        return redirect(url_for('welcome'))
+    if request.method == 'GET':
+        return render_template("index.html")
+    username = request.form.get('username')
+    password = request.form.get('password')
+    user = User.get_user(username, password)
+    if not user:
+        return render_template("index.html", error="User not found")
+    session['userid'] = user.id
+    return redirect(url_for('welcome'))
+
+
+@app.route('/welcome', methods=['GET'])
+def welcome():
+    user = get_logged_in_user()
+    if user is None:
+        return redirect(url_for("index"))
+    return render_template('welcome.html', user=user)
+
+
+@app.route('/update_age', methods=['POST'])
+def update_age():
+    user = get_logged_in_user()
+    if user is None:
+        return jsonify(status="User not logged in")
+    age = int(request.form.get('age'))
+    user.age = age
+    success = user.save()
+    if success[0]:
+        return jsonify(status="Age successfully changed")
+    else:
+        return jsonify(status="error " + success[1])
+
+
+@app.route('/logout')
+@csrf.exempt
+def logout():
+    if 'userid' in session:
+        session['userid'] = None
+    return redirect(url_for('index'))
+
+
+if __name__ == '__main__':
+    app.run(port=5000, debug=True)

web/login/models/__init__.py

@@ -0,0 +1 @@
+from user import *

web/login/models/user.py

@@ -0,0 +1,75 @@
+from mysql_db import MySQLdb
+
+import traceback
+
+
+class User:
+
+    def __init__(self, details):
+        self.id = details[0]
+        self.username = details[1]
+        self.name = details[3]
+        self.age = details[4]
+
+    def save(self):
+        db = MySQLdb()
+        cursor = db.conn.cursor()
+        rows = None
+        try:
+            user_input = (self.age, self.id)  # pass the input in the form of a tuple
+            query = "UPDATE tbl_users SET age = %s WHERE id = %s"
+            cursor.execute(query, user_input)
+            db.conn.commit()
+            cursor.close()
+            if cursor.rowcount != 0:
+                return True, cursor.rowcount
+            return False, cursor.rowcount
+
+        except Exception as e:
+            print(e)
+            traceback.print_exc()
+            cursor.close()
+            db.conn.close()
+            return False, e
+
+    @staticmethod
+    def get_user(username, password):
+        # creating the cursor
+        db = MySQLdb()
+        cursor = db.conn.cursor()
+        rows = None
+        try:
+            user_input = (username, password)  # pass the input in the form of a tuple
+            query = "SELECT * FROM tbl_users WHERE username = %s AND password = %s"
+            cursor.execute(query, user_input)
+            rows = cursor.fetchall()
+            cursor.close()
+            return User(rows[0])
+
+        except Exception as e:
+            print(e)
+            traceback.print_exc()
+            cursor.close()
+            db.conn.close()
+            return None
+
+    @staticmethod
+    def logged_user(user_id):
+        # creating the cursor
+        db = MySQLdb()
+        cursor = db.conn.cursor()
+        rows = None
+        try:
+            user_input = (user_id,)  # pass the input in the form of a tuple
+            query = "SELECT * FROM tbl_users WHERE id = %s"
+            cursor.execute(query, user_input)
+            rows = cursor.fetchall()
+            cursor.close()
+            return User(rows[0])
+
+        except Exception as e:
+            print(e)
+            traceback.print_exc()
+            cursor.close()
+            db.conn.close()
+            return None

web/login/static/js/change_age.js

@@ -0,0 +1,4 @@
+function changeAgeRequest() {
+    var form = document.getElementById("change_age_form");
+    alert(form);
+}

web/login/template/index.html

@@ -0,0 +1,15 @@
+<html>
+<head>
+    <title>Login</title>
+</head>
+<body>
+<form method="post">
+    Username: <input type="text" name="username"/> <br/>
+    Password: <input type="password" name="password"/> <br/>
+    <input type="submit"/> <br/>
+    {% if error is defined %}
+    Error: {{ error }}<br/>
+    {% endif %}
+</form>
+</body>
+</html>

web/login/template/welcome.html

@@ -0,0 +1,19 @@
+<html>
+<head>
+    <title>Login</title>
+    <script src="{{ url_for('static', filename='js/change_age.js') }}"></script>
+</head>
+<body>
+Hello {{ user.name }}, your age is {{ user.age }} <br/>
+<a href="/logout">Log out</a><br/>
+
+{% if notice is defined %}
+Notice: {{ notice }}<br/>
+{% endif %}
+
+<form id="change_age_form">
+    Change age: <input type="text" name="age"/><br/>
+    <input type="button" value="submit" onclick="changeAgeRequest()"/>
+</form>
+</body>
+</html>