add CSRF protection

app/main.py

@@ -3,7 +3,10 @@ import auth
 import db
 import re
 from os import environ, urandom
-
+from flask_wtf import FlaskForm
+from wtforms import StringField, PasswordField, BooleanField, EmailField, HiddenField
+from wtforms.validators import DataRequired
+from flask_wtf.csrf import CSRFProtect
 
 HOSTS_PATH = ( environ["HOSTS"] if "HOSTS" in environ else"/app/data/hosts.d/" ) + "/izbi"
 DB_PATH = environ["DB_PATH"] if "DB_PATH" in environ else "/app/data/sqlite.db"
@@ -14,8 +17,9 @@ app.config["SESSION_TYPE"] = "memcached"
 app.config["PERMANENT_SESSION_LIFETIME"] = 3600
 app.config["SECRET_KEY"] = environ["SECRET_KEY"] if "SECRET_KEY" in environ else urandom(32).hex()
 
-hosts = {}
+csrf = CSRFProtect(app)
 
+hosts = {}
 
 # Update hosts file
 def update_hosts(path: str, hosts: dict = {}) -> bool:
@@ -38,34 +42,61 @@ def validate_ip(ip: str) -> bool:
     return re.search(exp, ip)
 
 
+class LoginForm(FlaskForm):
+    username = StringField('username', validators=[DataRequired()])
+    password = PasswordField('password', validators=[DataRequired()])
+
 # Serve homepage
-@app.route('/', methods=["GET"])
-def homepage():
+@app.route('/', methods=["GET", "POST"])
+def index():
     if session.get("USERID") is not None:
         return redirect("/dashboard")
-    return render_template("index.html")
+
+    login_form = LoginForm()
+
+    if login_form.validate_on_submit():
+        username = login_form.username.data
+        password = login_form.password.data
+
+        userid = auth.user_login(DB_PATH, username, password)
+
+        if userid is None:
+            flash("Invalid user or password")
+            return redirect("/")
+
+        session["USERID"] = userid
+
+        return redirect("/dashboard")
+
+    return render_template("index.html", login_form=login_form)
 
 
-# Handle login
-@app.route('/login', methods=["POST"])
-def login():
-    userid = auth.user_login(DB_PATH, request.form.get("user"), request.form.get("pass"))
+class LogoutForm(FlaskForm):
+    pass
 
-    if userid is None:
-        flash("Invalid user or password")
-        return redirect("/")
+class GenerateTokenForm(FlaskForm):
+    token_name = StringField('token-name', validators=[DataRequired()])
 
-    session["USERID"] = userid
+class RevokeTokenForm(FlaskForm):
+    token_name = StringField('token-name', validators=[DataRequired()])
 
-    return redirect("/dashboard")
+class ChangePasswordForm(FlaskForm):
+    password = PasswordField('password', validators=[DataRequired()])
+    new_password = PasswordField('new-password', validators=[DataRequired()])
+    repeat_password = PasswordField('repeat-password', validators=[DataRequired()])
 
+class CreateUserForm(FlaskForm):
+    username = StringField('username', validators=[DataRequired()])
+    domains = StringField('domains')
+    password = PasswordField('password', validators=[DataRequired()])
+    is_admin = BooleanField('is-admin')
+    email = EmailField('email')
 
-# Log out by clearing session data
-@app.route('/logout', methods=["GET"])
-def logout():
-    session.clear()
-    return redirect("/")
+class UpdateUserForm(FlaskForm): #TODO
+    pass
 
+class DeleteUserForm(FlaskForm):
+    userid = HiddenField('userid', validators=[DataRequired()])
 
 # Display dashboard for logged in users
 @app.route('/dashboard', methods=["GET"])
@@ -73,6 +104,11 @@ def dashboard():
     if session.get("USERID") is None:
         return redirect("/")
 
+    logout_form = LogoutForm()
+    generate_token_form = GenerateTokenForm()
+    revoke_token_form = RevokeTokenForm()
+    change_password_form = ChangePasswordForm()
+
     userid = session.get("USERID")
 
     username = db.User.get_username(DB_PATH, userid)
@@ -81,10 +117,32 @@ def dashboard():
 
     # If user is admin, display additional options
     if db.User.is_admin(DB_PATH, userid):
-        users = db.list_users(DB_PATH)
-        return render_template("dashboard.html", username=username, domains=domains, tokens=tokens, is_admin=True, users=users)
+        create_user_form = CreateUserForm()
+        update_user_form = UpdateUserForm()
+        delete_user_form = DeleteUserForm()
 
-    return render_template("dashboard.html", username=username, domains=domains, tokens=tokens, is_admin=False)
+        users = db.list_users(DB_PATH)
+
+        return render_template("dashboard.html",
+            logout_form=logout_form, generate_token_form=generate_token_form, revoke_token_form=revoke_token_form, change_password_form=change_password_form, create_user_form=create_user_form, update_user_form=update_user_form, delete_user_form=delete_user_form,
+            username=username, domains=domains, tokens=tokens, is_admin=True, users=users
+        )
+
+    return render_template("dashboard.html",
+        logout_form=logout_form, generate_token_form=generate_token_form, revoke_token_form=revoke_token_form, change_password_form=change_password_form,
+        username=username, domains=domains, tokens=tokens, is_admin=False
+    )
+
+
+# Log out by clearing session data
+@app.route('/logout', methods=["POST"])
+def logout():
+    logout_form = LogoutForm()
+
+    if logout_form.validate_on_submit():
+        session.clear()
+
+    return redirect("/")
 
 
 # Generate new access token with a name
@@ -93,18 +151,18 @@ def generate_token():
     if session.get("USERID") is None:
         return redirect("/")
 
-    if "token" not in request.form:
-        flash("Token name not supplied")
-        return redirect("/dashboard")
+    generate_token_form = GenerateTokenForm()
 
-    try:
-        token = db.User.generate_token(DB_PATH, session.get("USERID"), request.form["token"])
-    except ValueError as e:
-        flash(str(e))
-        token = None
+    if generate_token_form.validate_on_submit():
+        token_name = generate_token_form.token_name.data
+        try:
+            token = db.User.generate_token(DB_PATH, session.get("USERID"), token_name)
+        except ValueError as e:
+            flash(str(e))
+            token = None
 
-    if token is not None:
-        flash(token)
+        if token is not None:
+            flash(token)
 
     return redirect("/dashboard")
 
@@ -115,14 +173,42 @@ def revoke_token():
     if session.get("USERID") is None:
         return redirect("/")
 
-    if "token" not in request.form:
-        flash("Token name not supplied")
-        return redirect("/dashboard")
+    revoke_token_form = RevokeTokenForm()
 
-    try:
-        db.User.revoke_token(DB_PATH, session.get("USERID"), request.form["token"])
-    except ValueError as e:
-        flash(str(e))
+    if revoke_token_form.validate_on_submit():
+        token_name = revoke_token_form.token_name.data
+        try:
+            db.User.revoke_token(DB_PATH, session.get("USERID"), token_name)
+        except ValueError as e:
+            flash(str(e))
+
+    return redirect("/dashboard")
+
+
+# Change password (by user)
+@app.route('/change-password', methods=["POST"])
+def change_password():
+    if session.get("USERID") is None:
+        return redirect("/")
+
+    change_password_form = ChangePasswordForm()
+
+    if change_password_form.validate_on_submit():
+        password = change_password_form.password.data
+        new_password = change_password_form.new_password.data
+        repeat_password = change_password_form.repeat_password.data
+
+        if new_password != repeat_password:
+            flash("Passwords do not match")
+            return redirect("/dashboard")
+
+        try:
+            auth.change_user_password(DB_PATH, session.get("USERID"), password, new_password)
+        except ValueError as e:
+            flash(str(e))
+            return redirect("/dashboard")
+
+        flash("Password changed successfully")
 
     return redirect("/dashboard")
 
@@ -135,37 +221,32 @@ def create_user():
     if not db.User.is_admin(DB_PATH, session.get("USERID")):
         return redirect("/dashboard")
 
-    if 'username' not in request.form or 'password' not in request.form:
-        flash('Required fields not supplied')
-        return redirect("/dashboard")
+    create_user_form = CreateUserForm()
 
-    username = request.form.get("username").strip()
-    password = request.form.get("password")
+    if create_user_form.validate_on_submit():
+        username = create_user_form.username.data.strip()
+        password = create_user_form.password.data
 
-    domains = request.form.get('domains')
-    if domains is None:
-        domains = []
-    else:
-        # Collapse whitespaces
-        collapse_exp = re.compile(r'\s+')
-        domains = collapse_exp.sub(' ', domains.strip())
-        domains = domains.split(" ")
+        domains = create_user_form.domains.data
+        if domains is None:
+            domains = []
+        else:
+            # Collapse whitespaces
+            collapse_exp = re.compile(r'\s+')
+            domains = collapse_exp.sub(' ', domains.strip())
+            domains = domains.split(" ")
 
-    is_admin = request.form.get("is_admin")
-    if is_admin is not None:
-        is_admin = is_admin == "true"
-    else:
-        is_admin = False
+        is_admin = create_user_form.is_admin.data
 
-    email = request.form.get("email")
-    if email is not None:
-        if email == "":
-            email = None
+        email = create_user_form.email.data
+        if email is not None:
+            if email == "":
+                email = None
 
-    try:
-        db.User.new(DB_PATH, username, password, domains, is_admin, email)
-    except ValueError as e:
-        flash(str(e))
+        try:
+            db.User.new(DB_PATH, username, password, domains, is_admin, email)
+        except ValueError as e:
+            flash(str(e))
 
     return redirect("/dashboard")
 
@@ -178,6 +259,11 @@ def update_user():
     if not db.User.is_admin(DB_PATH, session.get("USERID")):
         return redirect("/dashboard")
 
+    update_user_form = UpdateUserForm()
+
+    if not update_user_form.validate_on_submit():
+        return redirect("/dashboard")
+
     userid = request.form.get("userid")
     if userid is None:
         flash("UserID not supplied")
@@ -226,46 +312,20 @@ def delete_user():
     if not db.User.is_admin(DB_PATH, session.get("USERID")):
         return redirect("/dashboard")
 
-    userid = request.form.get("userid")
-    if userid is None:
-        return redirect("/dashboard")
+    delete_user_form = DeleteUserForm()
 
-    if db.User.exists(DB_PATH, userid):
-        flash("User does not exst")
-        return redirect("/dashboard")
+    if delete_user_form.validate_on_submit():
+        userid = delete_user_form.userid.data
 
-    db.User.delete(DB_PATH, userid)
+        if not db.User.exists(DB_PATH, userid):
+            flash("User does not exst")
+            return redirect("/dashboard")
+
+        db.User.delete(DB_PATH, userid)
 
     return redirect("/dashboard")
 
 
-# Change password (by user)
-@app.route('/change-password', methods=["POST"])
-def change_password():
-    if session.get("USERID") is None:
-        return redirect("/")
-
-    if "pass" not in request.form or "pass-new" not in request.form or "pass-rep" not in request.form:
-        flash("Required values not supplied")
-        return redirect("/dashboard")
-
-    oldpass = request.form["pass"]
-    newpass = request.form["pass-new"]
-    reppass = request.form["pass-rep"]
-
-    if newpass != reppass:
-        flash("Passwords do not match")
-        return redirect("/dashboard")
-
-    try:
-        auth.change_user_password(DB_PATH, session.get("USERID"), oldpass, newpass)
-    except ValueError as e:
-        flash(str(e))
-        return redirect("/dashboard")
-
-    flash("Password changed successfully")
-    return redirect("/dashboard")
-
 
 # API health endpoint ???
 @app.route('/health', methods=["GET"])
@@ -276,7 +336,8 @@ def health():
 # API for updating domain's IP
 @app.route('/update', methods=["GET"])
 def update_addr():
-    try:
+#    try:
+    if True:
         if 'domain' not in request.args:
             return jsonify({"status": "400", "code": "domain-error", "comment": "Missing or invalid domain"}), 400
         domain = request.args['domain']
@@ -307,9 +368,9 @@ def update_addr():
             return jsonify({"status": "500", "code": "internal-server-error", "comment": "An error occurred while processing request"}), 500
 
         return jsonify({"status": "200", "code": "ok", "comment": "Updated successfully"})
-    except Exception as e:
-        print(e, flush=True)
-        return jsonify({"status": "500", "code": "internal-server-error", "comment": "An error occurred while processing request"}), 500
+#    except Exception as e:
+#        print(e, flush=True)
+#        return jsonify({"status": "500", "code": "internal-server-error", "comment": "An error occurred while processing request"}), 500
 
 
 if __name__ == '__main__':