Add login/register system and database and also offensive mode.

app.py

@@ -1,10 +1,10 @@
-from flask import Flask, render_template, request
+from flask import Flask, render_template, request, g, redirect, url_for, Response
 from dotenv import load_dotenv
 from google.genai import Client, types
 
-from constants import SCENARIO_PROMPT, ANSWER_PROMPT, debt_amount_regex, evaluation_regex, NAME
+from constants import OFFENSIVE_SCENARIO_PROMPT, OFFENSIVE_ANSWER_PROMPT, debt_amount_regex, evaluation_regex, AI_NAME
 
-import os, requests, time, re
+import os, requests, time, re, sqlite3, flask_login, bcrypt, secrets
 
 load_dotenv(".env")
 
@@ -12,10 +12,123 @@ if not os.environ["USE_HACKCLUB_AI"]:
     gemini_client = Client()
 
 app = Flask(__name__)
+app.secret_key = os.environ["FLASK_SECRET_KEY"]
+
+login_manager = flask_login.LoginManager()
+login_manager.init_app(app)
+
+def get_db():
+    db = getattr(g, '_database', None)
+
+    if db is None:
+        db = g._database = sqlite3.connect(os.environ.get("DB_FILE", "data.db"))
+        db.execute("""
+            CREATE TABLE IF NOT EXISTS Users (
+                username TEXT PRIMARY KEY,
+                offended_debt_amount INT NOT NULL,
+                defended_debt_amount INT NOT NULL,
+                defensive_wins INT NOT NULL,
+                offensive_wins INT NOT NULL,
+                password TEXT NOT NULL,
+                password_salt TEXT NOT NULL
+            )
+        """)
+
+    return db
+
+@app.teardown_appcontext
+def close_connection(exception):
+    db = getattr(g, '_database', None)
+    if db is not None:
+        db.close()
+
+class User(flask_login.UserMixin):
+    pass
+
+@login_manager.user_loader
+def user_loader(user_id):
+    user = User()
+    user.id = user_id
+    return user
+
+@login_manager.unauthorized_handler
+def unathorized_handler():
+    return redirect(url_for("login"))
 
 @app.route("/")
+@flask_login.login_required
 def main():
-    return render_template("index.jinja2", name=NAME)
+    username = flask_login.current_user.id
+    return render_template("index.jinja2", username=username)
+
+@app.route("/offensive")
+@flask_login.login_required
+def offensive_mode():
+    username = flask_login.current_user.id
+    return render_template("offensive.jinja2", ai_name=AI_NAME, username=username)
+
+@app.route("/leaderboard")
+def leaderboard():
+    return render_template("leaderboard.jinja2")
+
+@app.route("/login", methods=["GET", "POST"])
+def login():
+    if hasattr(flask_login.current_user, "id"):
+        return redirect(url_for("main"))
+
+    if request.method == "GET":
+        return render_template("login.jinja2")
+    elif request.method == "POST":
+        username, password = request.form.get("username"), request.form.get("password")
+
+        cur = get_db().cursor()
+
+        cur.execute("SELECT password, password_salt FROM Users WHERE username = ?", (username,))
+
+        row = cur.fetchone()
+        if not row:
+            cur.close()
+            return redirect(url_for("login"))
+        
+        hashed_password, salt = row
+
+        if secrets.compare_digest(bcrypt.hashpw(password.encode(), salt.encode()), hashed_password.encode()):
+            cur.close()
+
+            user = User()
+            user.id = username
+            flask_login.login_user(user, remember=True)
+
+            return redirect(url_for("main"))
+        else:
+            cur.close()
+            return Response("Unathorized", 401)
+
+@app.route("/register", methods=["GET", "POST"])
+def register():
+    if hasattr(flask_login.current_user, "id"):
+        return redirect(url_for("main"))
+
+    if request.method == "GET":
+        return render_template("register.jinja2")
+    elif request.method == "POST":
+        username, password = request.form.get("username"), request.form.get("password")
+
+        cur = get_db().cursor()
+
+        cur.execute("SELECT username from Users WHERE username = ?", (username,))
+
+        if cur.fetchone():
+            return Response("An Account with this username already exists.", 400)
+        
+        salt = bcrypt.gensalt()
+        hashed_password = bcrypt.hashpw(password.encode(), salt)
+        
+        cur.execute("INSERT INTO Users (username, password, password_salt, offended_debt_amount, defended_debt_amount, defensive_wins, offensive_wins) VALUES (?, ?, ?, ?, ?, ?, ?)", (username, hashed_password.decode(), salt.decode(), 0, 0, 0, 0))
+        get_db().commit()
+        cur.close()
+
+        return redirect(url_for("login"))
 
 def ai_prompt(prompt):
     if os.environ["USE_HACKCLUB_AI"]:
@@ -36,12 +149,13 @@ def ai_prompt(prompt):
 
         return response.text.replace("'''", '')
 
-@app.route("/generate_scenario")
-def generate_scenario():
+@app.route("/offensive_scenario")
+@flask_login.login_required
+def offensive_scenario():
     text = ""
 
     while not "Debt amount: " in text or not "Scenario: " in text or not re.findall(debt_amount_regex, text):
-        text = ai_prompt(SCENARIO_PROMPT)
+        text = ai_prompt(OFFENSIVE_SCENARIO_PROMPT)
 
         time.sleep(0.5)
 
@@ -50,8 +164,9 @@ def generate_scenario():
         "debt_amount": int(text.split("Debt amount: ")[1].split("$")[0])
     }
 
-@app.route("/get_answer", methods=["POST"])
-def get_answer():
+@app.route("/offensive_answer", methods=["POST"])
+@flask_login.login_required
+def offensive_answer():
     scenario, user_input = request.json['scenario'], request.json["user_input"]
 
     if not scenario or not user_input:
@@ -60,15 +175,13 @@ def get_answer():
     text = ""
 
     while not re.findall(evaluation_regex, text):
-        text = ai_prompt(ANSWER_PROMPT.format_map({"scenario": scenario, "user_input": user_input, "name": NAME}))
+        text = ai_prompt(OFFENSIVE_ANSWER_PROMPT.format_map({"scenario": scenario, "user_input": user_input, "ai_name": AI_NAME}))
 
         time.sleep(0.5)
 
-    print(text.split("Convinced: "), text.split("Convinced: ")[1])
-
     return {
         "story": text.split("\nEVALUATION")[0],
-        "convinced": True if text.split("Convinced: ")[1].split("\nFinal")[0] == "Yes" else False,
+        "convinced": True if "Yes" in text.split("Convinced: ")[1].split("\nFinal")[0] else False,
         "final_debt_amount": text.split("Final Debt Amount: ")[1].split("$")[0]
     }