# Demonstrates session and cookie processing. The username is stored
# as a cookie and an internal userid is saved in a session variable.
# Also demonstrates secure password authentication using bcrypt salt
# and hash.
#
# References:
#   https://en.wikibooks.org/wiki/Python_Programming
#   https://flask.palletsprojects.com/en/1.1.x/quickstart/
#   http://zetcode.com/python/bcrypt/

import os
import flask
import bcrypt

app = flask.Flask(__name__)
app.secret_key = os.urandom(32) # 256-bit key secures session variables

FORM_NO_COOKIE = """
<form method="POST">
<p>Please log in:</p>
<p><label for="username">Username:</label>
<input type="text" id="username" name="username"></p>
<p><label for="password">Password:</label>
<input type="password" id="password" name="password"></p>
<input type="submit" id="log-in" name="log-in" value="Log In">
</form>
"""

FORM_COOKIE_1 = """
<form method="POST">
<p>Please log in:</p>
<p><label for="username">Username:</label>
<input type="text" id="username" name="username" 
"""

FORM_COOKIE_2 = """
></p>
<p><label for="password">Password:</label>
<input type="password" id="password" name="password"></p>
<input type="submit" id="log-in" name="log-in" value="Log In">
<input type="submit" id="forget-me" name="forget-me" value="Forget Me">
</form>
"""

FORM_LOGOUT = """
<form method="POST">
<input type="submit" id="log-out" name="log-out" value="Log Out">
</form>
"""

users = [
    # Password is the same as the username, just salted and hashed.
    # Don't do this in a production application! Use custom passwords.
    { "userid": 1, "username": "admin", 
        "password": b'$2b$12$6xEcJ9bCRo3JgNWyn32fwuSoRh1pg8f81jjHpYq6NQ9Y8uDkhWOE6'},
    { "userid": 2, "username": "test", 
        "password": b'$2b$12$UZLEFMg9ez.n88Sjpb/ZN.VVlmyPPxHOeL/DE452Si4H3PSQSB0Pa'}
]

@app.route('/', methods=["GET"])
def root_get():
    userid = flask.session.get("userid")
    username = flask.request.cookies.get("username")
    result = f"<p>Session: {bool(userid)}. Cookie: {bool(username)}.</p>"

    if userid:
        # Already logged in.
        return result + f"<p>Welcome back {username}!</p>" + FORM_LOGOUT
    elif username:
        # Returning user.
        return result + FORM_COOKIE_1 + f"value=\"{username}\"" + FORM_COOKIE_2
    else:
        return result + FORM_NO_COOKIE
 
@app.route('/', methods=["POST"])
def root_post():
    if flask.request.form.get('log-out', None):
        username = flask.request.cookies.get("username")
        flask.session.clear()
        result = f"<p>Session: False. Cookie: {bool(username)}.</p>" + \
            f"<p>Logged out {username}!</p>"
        return result

    elif flask.request.form.get('forget-me', None):
        flask.session.clear()
        result = "<p>Session: False. Cookie: False.</p>" + \
            "<p>Good bye!</p>"
        response = flask.make_response(result)
        response.set_cookie("username", "", max_age=0)
        return response

    else:
        username = flask.request.form["username"]
        password = flask.request.form["password"]
        userid = authenticate_user(username, password)
        if userid:
            flask.session["userid"] = userid
            result = "<p>Session: True. Cookie: True.</p>" + \
                f"<p>Welcome {username}!</p>"
            response = flask.make_response(result)
            response.set_cookie("username", username)
            return response
        else:
            return flask.redirect("/", code=303)

def authenticate_user(username, password):
    for user in users:
        if user["username"] == username:
            result = bcrypt.checkpw(password.encode(), user["password"])
            if result:
                return user["userid"]
            else:
                # Should track number of failed attempts, lock account, etc.
                return None
    return None

def generate_hashed_password(password):
    # Use this function to generate hashed passwords to save in 
    # the users list or a database.
    salt = bcrypt.gensalt()
    hashed = bcrypt.hashpw(password, salt)
    return hashed

if __name__ == "__main__":
    app.run(host='0.0.0.0', port=5000)