login.py

from main import app
from flask import Flask, redirect, url_for, render_template, abort, request, make_response, session, flash
import sqlite3

# User directed to login as default
@app.route('/')
def redirect_to():
    return redirect(url_for('login'))

@app.route('/login',methods=['GET','POST'])
def login():
    # Recieve user_name and passwrd values from form in template
    if request.method == 'POST':
        return do_login(request.form['user_name'], request.form['passwrd'])
    else:
        # If either access modes have been stored before, login not required
        if "access_mode" in session:
            return redirect(url_for('homepage'))
        return render_template('login.html', page=url_for('login'), register_page=url_for('register'))
    

def do_login(user_name,passwrd):
    # Simple DB stuff learnt in 4007CEM 
    con = sqlite3.connect('bookshop_database.db')
    # Creates list, easier to work with than tuples (sqlite3)
    con.row_factory = lambda cursor, row: row[0]
    cur = con.cursor()
    # Attempt to match login with database (return either 0 or 1 rows)
    cur.execute("SELECT password FROM users WHERE username=? AND password=?",(user_name,passwrd))
    # Fetch all rows into variable
    row = cur.fetchall()

    # IF login exists
    if len(row) != 0:
        if user_name == 'admin':
            # Closing DB good practice. Memory issues 
            con.close()
            return go_to_homepage('admin')
        else:
            con.close()
            return go_to_homepage('user')
    else:
        con.close()
        # Simple login failed message to template
        flash('Incorrect login/password')
        return redirect(url_for('login'))

# This function used to communicate the users permissions to template
def go_to_homepage(access_mode):
    session["access_mode"] = access_mode
    return redirect(url_for('homepage'))

# Simple register form with little input validation, could cause issues as users can in theory create infinite users
@app.route('/register',methods=['GET','POST'])
def register():
    if request.method =='POST':
        return do_register(request.form['user_name'], request.form['passwrd'])
    else:
        return render_template('register.html',page=url_for('register'))
    

def do_register(user_name, passwrd):
    # HTML checks that both username and pass must be a certain length
    # check there is no ' ' (whitespace)
    if (' ' in passwrd) == True or (' ' in user_name) == True:
        flash('Do not put whitespaces in username or password!')
        # Restart if any ' ' found (whitespace)
        return redirect(url_for('register'))
    # Connect to DB etc and compare username to existing users
    con = sqlite3.connect('bookshop_database.db')
    con.row_factory = lambda cursor, row: row[0]
    cur = con.cursor()
    cur.execute("SELECT username FROM users WHERE username=?",(user_name,))
    row = cur.fetchall() 
    # If the user doesnt exist we can insert into DB        
    if len(row) == 0:
        cur.execute("INSERT INTO users VALUES(?,?)",(user_name,passwrd))
        con.commit()
        con.close()
        return redirect(url_for('login'))
    # If username exists, start form again. 
    else:
        con.close()
        flash('Username already taken!!!')
        return redirect(url_for('register'))