quick cleanup? (untested)

src/main.py

@@ -6,33 +6,9 @@ import secrets
 import os
 from werkzeug.utils import secure_filename
 
-
 UPLOAD_FOLDER = "static/uploads"
 os.makedirs(UPLOAD_FOLDER, exist_ok=True)
 
-COUNTRIES = [
-    "Afghanistan","Albania","Algeria","Andorra","Angola","Antigua and Barbuda","Argentina","Armenia","Australia","Austria",
-    "Azerbaijan","Bahamas","Bahrain","Bangladesh","Barbados","Belarus","Belgium","Belize","Benin","Bhutan",
-    "Bolivia","Bosnia and Herzegovina","Botswana","Brazil","Brunei","Bulgaria","Burkina Faso","Burundi","Cabo Verde","Cambodia",
-    "Cameroon","Canada","Central African Republic","Chad","Chile","China","Colombia","Comoros","Congo (Congo-Brazzaville)","Costa Rica",
-    "Croatia","Cuba","Cyprus","Czechia (Czech Republic)","Democratic Republic of the Congo","Denmark","Djibouti","Dominica","Dominican Republic","Ecuador",
-    "Egypt","El Salvador","Equatorial Guinea","Eritrea","Estonia","Eswatini (fmr. Swaziland)","Ethiopia","Fiji","Finland","France",
-    "Gabon","Gambia","Georgia","Germany","Ghana","Greece","Grenada","Guatemala","Guinea","Guinea-Bissau",
-    "Guyana","Haiti","Holy See","Honduras","Hungary","Iceland","India","Indonesia","Iran","Iraq",
-    "Ireland","Israel","Italy","Jamaica","Japan","Jordan","Kazakhstan","Kenya","Kiribati","Kuwait",
-    "Kyrgyzstan","Laos","Latvia","Lebanon","Lesotho","Liberia","Libya","Liechtenstein","Lithuania","Luxembourg",
-    "Madagascar","Malawi","Malaysia","Maldives","Mali","Malta","Marshall Islands","Mauritania","Mauritius","Mexico",
-    "Micronesia","Moldova","Monaco","Mongolia","Montenegro","Morocco","Mozambique","Myanmar (Burma)","Namibia","Nauru",
-    "Nepal","Netherlands","New Zealand","Nicaragua","Niger","Nigeria","North Korea","North Macedonia","Norway","Oman",
-    "Pakistan","Palau","Palestine State","Panama","Papua New Guinea","Paraguay","Peru","Philippines","Poland","Portugal",
-    "Qatar","Romania","Russia","Rwanda","Saint Kitts and Nevis","Saint Lucia","Saint Vincent and the Grenadines","Samoa","San Marino","Sao Tome and Principe",
-    "Saudi Arabia","Senegal","Serbia","Seychelles","Sierra Leone","Singapore","Slovakia","Slovenia","Solomon Islands","Somalia",
-    "South Africa","South Korea","South Sudan","Spain","Sri Lanka","Sudan","Suriname","Sweden","Switzerland","Syria",
-    "Tajikistan","Tanzania","Thailand","Timor-Leste","Togo","Tonga","Trinidad and Tobago","Tunisia","Turkey","Turkmenistan",
-    "Tuvalu","Uganda","Ukraine","United Arab Emirates","United Kingdom","United States","Uruguay","Uzbekistan","Vanuatu","Venezuela",
-    "Vietnam","Yemen","Zambia","Zimbabwe"
-]
-
 app = Flask(__name__)
 app.config['SQLALCHEMY_DATABASE_URI'] = 'mysql+pymysql://love:love@localhost:3309/lovedb'
 app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
@@ -42,93 +18,121 @@ app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
 db = SQLAlchemy(app)
 gpg = gnupg.GPG()
 
+COUNTRIES = [ "Afghanistan","Albania","Algeria","Andorra","Angola","Antigua and Barbuda","Argentina",
+              "Armenia","Australia","Austria","Azerbaijan","Bahamas","Bahrain","Bangladesh",
+              "Barbados","Belarus","Belgium","Belize","Benin","Bhutan","Bolivia","Bosnia and Herzegovina",
+              "Botswana","Brazil","Brunei","Bulgaria","Burkina Faso","Burundi","Cabo Verde","Cambodia",
+              "Cameroon","Canada","Central African Republic","Chad","Chile","China","Colombia","Comoros",
+              "Congo (Congo-Brazzaville)","Costa Rica","Croatia","Cuba","Cyprus","Czechia (Czech Republic)",
+              "Democratic Republic of the Congo","Denmark","Djibouti","Dominica","Dominican Republic","Ecuador",
+              "Egypt","El Salvador","Equatorial Guinea","Eritrea","Estonia","Eswatini (fmr. Swaziland)",
+              "Ethiopia","Fiji","Finland","France","Gabon","Gambia","Georgia","Germany","Ghana","Greece",
+              "Grenada","Guatemala","Guinea","Guinea-Bissau","Guyana","Haiti","Holy See","Honduras","Hungary",
+              "Iceland","India","Indonesia","Iran","Iraq","Ireland","Israel","Italy","Jamaica","Japan","Jordan",
+              "Kazakhstan","Kenya","Kiribati","Kuwait","Kyrgyzstan","Laos","Latvia","Lebanon","Lesotho",
+              "Liberia","Libya","Liechtenstein","Lithuania","Luxembourg","Madagascar","Malawi","Malaysia",
+              "Maldives","Mali","Malta","Marshall Islands","Mauritania","Mauritius","Mexico","Micronesia",
+              "Moldova","Monaco","Mongolia","Montenegro","Morocco","Mozambique","Myanmar (Burma)","Namibia",
+              "Nauru","Nepal","Netherlands","New Zealand","Nicaragua","Niger","Nigeria","North Korea",
+              "North Macedonia","Norway","Oman","Pakistan","Palau","Palestine State","Panama","Papua New Guinea",
+              "Paraguay","Peru","Philippines","Poland","Portugal","Qatar","Romania","Russia","Rwanda",
+              "Saint Kitts and Nevis","Saint Lucia","Saint Vincent and the Grenadines","Samoa","San Marino",
+              "Sao Tome and Principe","Saudi Arabia","Senegal","Serbia","Seychelles","Sierra Leone","Singapore",
+              "Slovakia","Slovenia","Solomon Islands","Somalia","South Africa","South Korea","South Sudan",
+              "Spain","Sri Lanka","Sudan","Suriname","Sweden","Switzerland","Syria","Tajikistan","Tanzania",
+              "Thailand","Timor-Leste","Togo","Tonga","Trinidad and Tobago","Tunisia","Turkey","Turkmenistan",
+              "Tuvalu","Uganda","Ukraine","United Arab Emirates","United Kingdom","United States","Uruguay",
+              "Uzbekistan","Vanuatu","Venezuela","Vietnam","Yemen","Zambia","Zimbabwe" ]
+
+
 class User(db.Model):
     id = db.Column(db.Integer, primary_key=True)
-
     username = db.Column(db.String(128), unique=True, nullable=False)
     pgp = db.Column(db.String(4096), nullable=False)
-
     firstname = db.Column(db.String(128), nullable=False)
     lastname = db.Column(db.String(128), nullable=False)
-
     sex = db.Column(db.Enum('male', 'female'), nullable=False)
     date_of_birth = db.Column(db.Date, nullable=False)
-
     profile_picture = db.Column(db.String(200), nullable=False)
-
     pictures = db.Column(db.JSON, nullable=True)
-
     country = db.Column(db.String(128), nullable=False)
     city = db.Column(db.String(128), nullable=True)
-
     height = db.Column(db.Float, nullable=True)
     weight = db.Column(db.Integer, nullable=True)
     race = db.Column(db.String(20), nullable=True)
-
     prefered_age_range = db.Column(db.String(20), nullable=True)
-
     likes = db.Column(db.JSON, nullable=True)
     dislikes = db.Column(db.JSON, nullable=True)
-
     xmpp = db.Column(db.String(128), unique=True, nullable=False)
     email = db.Column(db.String(128), unique=True, nullable=True)
     phone = db.Column(db.String(20), unique=True, nullable=True)
-
     is_verified = db.Column(db.Boolean, default=False)
 
+
+def calculate_age(dob: date) -> int:
+    today = date.today()
+    return today.year - dob.year - ((today.month, today.day) < (dob.month, dob.day))
+
+def save_files(username: str, profile_file, pictures_files):
+    user_folder = os.path.join(app.config['UPLOAD_FOLDER'], username)
+    os.makedirs(user_folder, exist_ok=True)
+
+    profile_filename = secure_filename(profile_file.filename)
+    profile_path = os.path.join(user_folder, profile_filename)
+    profile_file.save(profile_path)
+    profile_url = f"/{profile_path.replace(os.sep, '/')}"
+
+    pictures_urls = []
+    for pic in pictures_files:
+        if pic.filename:
+            filename = secure_filename(pic.filename)
+            path = os.path.join(user_folder, filename)
+            pic.save(path)
+            pictures_urls.append(f"/{path.replace(os.sep, '/')}")
+
+    return profile_url, pictures_urls
+
+def pgp_encrypt_and_import(pgp_key: str, message: str):
+    result = gpg.import_keys(pgp_key)
+    if not result.fingerprints:
+        return None, None
+    fingerprint = result.fingerprints[0]
+    encrypted = gpg.encrypt(message, recipients=[fingerprint])
+    if not encrypted.ok:
+        return fingerprint, None
+    return fingerprint, str(encrypted)
+
+
 @app.route("/")
 def home():
     return render_template("index.html")
 
+
 @app.route("/register", methods=["GET", "POST"])
 def register():
     if request.method == "POST":
-        username = request.form.get("username")
+        data = {key: request.form.get(key) for key in [
-        pgp = request.form.get("pgp")
+            "username","pgp","firstname","lastname","sex","date_of_birth","country","xmpp",
-        firstname = request.form.get("firstname")
+            "email","phone","city","height","weight","race","prefered_age_range"
-        lastname = request.form.get("lastname")
+        ]}
-        sex = request.form.get("sex")
-        date_of_birth = request.form.get("date_of_birth")
-        country = request.form.get("country")
-        xmpp = request.form.get("xmpp")
 
-        email = request.form.get("email")
+        required_fields = ["username","pgp","firstname","lastname","sex","date_of_birth","country","xmpp"]
-        phone = request.form.get("phone")
+        if not all(data[f] for f in required_fields):
-        city = request.form.get("city")
-        height = request.form.get("height")
-        weight = request.form.get("weight")
-        race = request.form.get("race")
-        prefered_age_range = request.form.get("prefered_age_range")
-
-        if not all([username, pgp, firstname, lastname, sex, date_of_birth, country, xmpp]):
             flash("Please fill all required fields.")
             return redirect(url_for("register"))
 
-        if User.query.filter_by(username=username).first():
+        for field in ["username","xmpp","email","phone"]:
-            flash("Username already exists.")
+            if data.get(field) and User.query.filter_by(**{field:data[field]}).first():
-            return redirect(url_for("register"))
+                flash(f"{field.capitalize()} already exists.")
-
-        if User.query.filter_by(xmpp=xmpp).first():
-            flash("XMPP already exists.")
-            return redirect(url_for("register"))
-
-        if email and User.query.filter_by(email=email).first():
-            flash("Email already exists.")
-            return redirect(url_for("register"))
-
-        if phone and User.query.filter_by(phone=phone).first():
-            flash("Phone already exists.")
                 return redirect(url_for("register"))
 
         try:
-            dob = date.fromisoformat(date_of_birth)
+            dob = date.fromisoformat(data["date_of_birth"])
         except ValueError:
             flash("Invalid date format.")
             return redirect(url_for("register"))
 
-        today = date.today()
+        if calculate_age(dob) < 18:
-        age = today.year - dob.year - ((today.month, today.day) < (dob.month, dob.day))
-        if age < 18:
             flash("You must be at least 18 years old to register.")
             return redirect(url_for("register"))
 
@@ -139,61 +143,24 @@ def register():
             flash("Profile picture is required.")
             return redirect(url_for("register"))
 
-        user_folder = os.path.join(app.root_path, "static/uploads", username)
+        profile_url, pictures_urls = save_files(data["username"], profile_file, pictures_files)
-        os.makedirs(user_folder, exist_ok=True)
 
-        from werkzeug.utils import secure_filename
-        profile_filename = secure_filename(profile_file.filename)
-        profile_path = os.path.join(user_folder, profile_filename)
-        profile_file.save(profile_path)
-        profile_url = f"/static/uploads/{username}/{profile_filename}"
-
-        pictures_urls = []
-        for pic in pictures_files:
-            if pic.filename:
-                filename = secure_filename(pic.filename)
-                path = os.path.join(user_folder, filename)
-                pic.save(path)
-                pictures_urls.append(f"/static/uploads/{username}/{filename}")
-
-        session["pending_user"] = {
-            "username": username,
-            "pgp": pgp,
-            "firstname": firstname,
-            "lastname": lastname,
-            "sex": sex,
-            "date_of_birth": date_of_birth,
-            "profile_url": profile_url,
-            "pictures_urls": pictures_urls,
-            "country": country,
-            "xmpp": xmpp,
-            "email": email,
-            "phone": phone,
-            "city": city,
-            "height": height,
-            "weight": weight,
-            "race": race,
-            "prefered_age_range": prefered_age_range
-        }
-
-        import_result = gpg.import_keys(pgp)
-        if not import_result.fingerprints:
-            flash("Invalid PGP key.")
-            return redirect(url_for("register"))
-
-        fingerprint = import_result.fingerprints[0]
         random_string = secrets.token_hex(16)
         challenge_phrase = f"this is the unencrypted string: {random_string}"
-        encrypted_data = gpg.encrypt(challenge_phrase, recipients=[fingerprint])
+        fingerprint, encrypted_msg = pgp_encrypt_and_import(data["pgp"], challenge_phrase)
-        if not encrypted_data.ok:
+
-            flash("Failed to encrypt challenge.")
+        if not fingerprint or not encrypted_msg:
+            flash("Invalid PGP key or encryption failed.")
             return redirect(url_for("register"))
 
+        session["pending_user"] = {**data, "profile_url": profile_url, "pictures_urls": pictures_urls}
         session["pgp_expected_phrase"] = challenge_phrase
-        return render_template("verify.html", encrypted_message=str(encrypted_data))
+
+        return render_template("verify.html", encrypted_message=encrypted_msg)
 
     return render_template("register.html", countries=COUNTRIES)
 
+
 @app.route("/verify", methods=["POST"])
 def verify():
     expected_phrase = session.get("pgp_expected_phrase")
@@ -208,9 +175,11 @@ def verify():
         flash("You must paste the decrypted message.")
         return redirect(url_for("register"))
 
-    if submitted.strip() == expected_phrase:
+    if submitted.strip() != expected_phrase:
-        dob = date.fromisoformat(data["date_of_birth"])
+        flash("Verification failed. Account not created.")
+        return redirect(url_for("register"))
 
+    dob = date.fromisoformat(data["date_of_birth"])
     new_user = User(
         username=data["username"],
         pgp=data["pgp"],
@@ -222,13 +191,13 @@ def verify():
         pictures=data["pictures_urls"],
         country=data["country"],
         xmpp=data["xmpp"],
-            email=data["email"] or None,
+        email=data.get("email") or None,
-            phone=data["phone"] or None,
+        phone=data.get("phone") or None,
-            city=data["city"] or None,
+        city=data.get("city") or None,
-            height=float(data["height"]) if data["height"] else None,
+        height=float(data["height"]) if data.get("height") else None,
-            weight=int(data["weight"]) if data["weight"] else None,
+        weight=int(data["weight"]) if data.get("weight") else None,
-            race=data["race"] or None,
+        race=data.get("race") or None,
-            prefered_age_range=data["prefered_age_range"] or None,
+        prefered_age_range=data.get("prefered_age_range") or None,
         is_verified=True
     )
 
@@ -237,25 +206,20 @@ def verify():
 
     session['user_id'] = new_user.id
     session['username'] = new_user.username
-
-        # limpar dados temporários
     session.pop("pending_user", None)
     session.pop("pgp_expected_phrase", None)
 
     flash("PGP verification successful! Account created.")
     return redirect(url_for("home"))
 
-    else:
-        flash("Verification failed. Account not created.")
-        return redirect(url_for("register"))
 
-@app.route("/login", methods=["GET", "POST"])
+@app.route("/login", methods=["GET","POST"])
 def login():
     if request.method == "POST":
         username = request.form.get("username")
-        pgp = request.form.get("pgp")
+        pgp_key = request.form.get("pgp")
 
-        if not username or not pgp:
+        if not username or not pgp_key:
             flash("Please enter both username and PGP key.")
             return redirect(url_for("login"))
 
@@ -264,35 +228,22 @@ def login():
             flash("User not found.")
             return redirect(url_for("login"))
 
-        import_result = gpg.import_keys(pgp)
-        if not import_result.fingerprints:
-            flash("Invalid PGP key.")
-            return redirect(url_for("login"))
-
-        fingerprint = import_result.fingerprints[0]
-
         random_string = secrets.token_hex(16)
         challenge_phrase = f"this is the unencrypted string: {random_string}"
+        fingerprint, encrypted_msg = pgp_encrypt_and_import(pgp_key, challenge_phrase)
 
-        encrypted_data = gpg.encrypt(
+        if not fingerprint or not encrypted_msg:
-            challenge_phrase,
+            flash("Invalid PGP key or encryption failed.")
-            recipients=[fingerprint]
-        )
-
-        if not encrypted_data.ok:
-            flash("Failed to encrypt challenge.")
             return redirect(url_for("login"))
 
         session["login_user_id"] = user.id
         session["login_expected_phrase"] = challenge_phrase
 
-        return render_template(
+        return render_template("login_verify.html", encrypted_message=encrypted_msg)
-            "login_verify.html",
-            encrypted_message=str(encrypted_data)
-        )
 
     return render_template("login.html")
 
+
 @app.route("/login_verify", methods=["POST"])
 def login_verify():
     user_id = session.get("login_user_id")
@@ -307,19 +258,19 @@ def login_verify():
         flash("You must paste the decrypted message")
         return redirect(url_for("login"))
 
-    if submitted.strip() == expected_phrase:
+    if submitted.strip() != expected_phrase:
+        flash("Verification failed")
+        return redirect(url_for("login"))
+
     user = User.query.get(user_id)
     session['user_id'] = user.id
     session['username'] = user.username
-
     session.pop("login_user_id", None)
     session.pop("login_expected_phrase", None)
 
     flash("Logged in successfully")
     return redirect(url_for("home"))
-    else:
+
-        flash("Verification failed")
-        return redirect(url_for("login"))
 
 @app.route("/logout")
 def logout():
@@ -328,12 +279,13 @@ def logout():
     flash("Logged out successfully")
     return redirect(url_for("home"))
 
+
 @app.route("/user/<username>")
 def user_profile(username):
     user = User.query.filter_by(username=username).first_or_404()
-
     return render_template("user.html", user=user, date=date)
 
+
 if __name__ == "__main__":
     with app.app_context():
         db.create_all()