Update server.py

server.py

@@ -1,104 +1,67 @@
 import os
-import sqlite3
-import secrets
+import csv
+import random
+import string
+import pathlib
 from datetime import datetime, timedelta
-from flask import Flask, request, redirect, url_for, render_template, session, flash
+from flask import Flask, request, redirect, url_for, session, render_template, flash, jsonify
 from flask_limiter import Limiter
 from flask_limiter.util import get_remote_address
-from huggingface_hub import HfApi
+from huggingface_hub import HfApi, HfFolder
+from datasets import Dataset as HFDataset
 import pandas as pd
 
-# ENV
+# --- Setup ---
+app = Flask(__name__)
+app.secret_key = os.getenv("FLASK_SECRET", "supersecret")
+limiter = Limiter(get_remote_address)
+limiter.init_app(app)
+
+# --- Admin credentials ---
 ADMIN_USER = os.getenv("ADMIN_USER", "admin")
 ADMIN_PASS = os.getenv("ADMIN_PASS", "Welcome123")
-FLASK_SECRET = os.getenv("FLASK_SECRET", "supersecret")
 HF_TOKEN = os.getenv("HF_TOKEN")
-HF_REPO_ID = "0vergeared/otp-logs"
-
-# Writeable database path
-DATABASE = "/tmp/otp.db"
-EXPORT_CSV = "/tmp/otp_all.csv"
-
-# Flask app
-app = Flask(__name__, template_folder="templates")
-app.secret_key = FLASK_SECRET
-limiter = Limiter(get_remote_address, app=app)
-
-# Initialize DB
-def init_db():
-    with sqlite3.connect(DATABASE) as conn:
-        c = conn.cursor()
-        c.execute('''CREATE TABLE IF NOT EXISTS otps (
-            id INTEGER PRIMARY KEY AUTOINCREMENT,
-            otp TEXT UNIQUE,
-            generated_at TEXT,
-            expires_at TEXT,
-            used_at TEXT,
-            ip_address TEXT,
-            status TEXT
-        )''')
-        conn.commit()
-
-# Generate OTP
-def generate_otp():
-    return ''.join(secrets.choice("0123456789") for _ in range(6))
-
-# Upload CSV to HF dataset
-def upload_to_hf(otp_row):
-    file_exists = os.path.exists(EXPORT_CSV)
-    df_new = pd.DataFrame([otp_row], columns=["otp", "generated_at", "expires_at", "used_at", "ip", "status"])
-
-    if file_exists:
-        df_old = pd.read_csv(EXPORT_CSV)
-        df = pd.concat([df_old, df_new], ignore_index=True)
-    else:
-        df = df_new
-
-    df.to_csv(EXPORT_CSV, index=False)
+
+# --- Dataset info ---
+REPO_ID = "0vergeared/otp-logs"
+LOCAL_CSV = "otp_temp.csv"
+
+# --- Ensure templates folder exists ---
+pathlib.Path("templates").mkdir(parents=True, exist_ok=True)
+
+# --- OTP Generator ---
+def generate_otp(length=6):
+    return ''.join(random.choices(string.digits, k=length))
+
+# --- Save OTP to Hugging Face Dataset ---
+def save_otp_to_dataset(otp: str, expiry: str):
+    if not HF_TOKEN:
+        print("❌ HF_TOKEN missing. Skipping dataset upload.")
+        return
 
     try:
+        HfFolder.save_token(HF_TOKEN)
         api = HfApi()
-        api.upload_file(
-            path_or_fileobj=EXPORT_CSV,
-            path_in_repo="otp_all.csv",
-            repo_id=HF_REPO_ID,
-            token=HF_TOKEN,
-            create_pr=False
-        )
-        print("✅ OTP log uploaded to HF.")
+
+        # Load existing data or create new
+        df = pd.read_csv(LOCAL_CSV) if os.path.exists(LOCAL_CSV) else pd.DataFrame(columns=["otp", "expiry", "used"])
+
+        # Append new OTP
+        df = pd.concat([df, pd.DataFrame([{"otp": otp, "expiry": expiry, "used": False}])], ignore_index=True)
+
+        # Save to local file
+        df.to_csv(LOCAL_CSV, index=False)
+
+        # Push to HF Hub
+        dataset = HFDataset.from_pandas(df)
+        dataset.push_to_hub(REPO_ID, token=HF_TOKEN)
+        print("✅ OTP pushed to Hugging Face dataset.")
     except Exception as e:
-        print("⚠️ Hugging Face upload failed:", e)
+        print("❌ Failed to save OTP to dataset:", e)
 
-# Routes
-@app.route("/", methods=["GET", "POST"])
+# --- Routes ---
+@app.route("/")
 def index():
-    if request.method == "POST":
-        otp_input = request.form.get("otp")
-        now = datetime.utcnow()
-
-        with sqlite3.connect(DATABASE) as conn:
-            c = conn.cursor()
-            c.execute("SELECT * FROM otps WHERE otp = ?", (otp_input,))
-            row = c.fetchone()
-
-            if row:
-                otp, gen_at, exp_at, used_at, ip, status = row[1:7]
-                exp_time = datetime.fromisoformat(exp_at)
-                if used_at:
-                    flash("OTP already used!", "danger")
-                elif now > exp_time:
-                    flash("OTP expired!", "danger")
-                else:
-                    used_time = now.isoformat()
-                    c.execute("UPDATE otps SET used_at = ?, status = ? WHERE otp = ?", (used_time, "used", otp_input))
-                    conn.commit()
-                    otp_row = [otp_input, gen_at, exp_at, used_time, request.remote_addr, "used"]
-                    upload_to_hf(otp_row)
-                    flash("✅ OTP accepted. Installation can proceed!", "success")
-                    return redirect(url_for("index"))
-            else:
-                flash("❌ Invalid OTP!", "danger")
-
     return render_template("index.html")
 
 @app.route("/admin", methods=["GET", "POST"])
@@ -109,38 +72,72 @@ def admin():
         if username == ADMIN_USER and password == ADMIN_PASS:
             session["admin"] = True
             return redirect(url_for("dashboard"))
-        else:
-            flash("❌ Login failed", "danger")
+        flash("❌ Login failed", "danger")
     return render_template("admin_login.html")
 
-@app.route("/dashboard", methods=["GET", "POST"])
+@app.route("/dashboard")
 def dashboard():
+    if not session.get("admin"):
+        return redirect(url_for("admin"))
+    otps = []
+    if os.path.exists(LOCAL_CSV):
+        with open(LOCAL_CSV, newline='') as f:
+            reader = csv.DictReader(f)
+            otps = list(reader)
+    return render_template("dashboard.html", otps=otps)
+
+@app.route("/generate_otp", methods=["POST"])
+def generate_otp_route():
     if not session.get("admin"):
         return redirect(url_for("admin"))
 
-    if request.method == "POST":
-        otp = generate_otp()
-        now = datetime.utcnow()
-        expires = now + timedelta(minutes=10)
-        ip = request.remote_addr
+    otp = generate_otp()
+    expiry = (datetime.utcnow() + timedelta(minutes=5)).isoformat()
+    save_otp_to_dataset(otp, expiry)
+    flash(f"✅ OTP generated: {otp}", "success")
+    return redirect(url_for("dashboard"))
+
+@app.route("/validate", methods=["POST"])
+@limiter.limit("10 per minute")
+def validate():
+    data = request.get_json()
+    input_otp = data.get("otp")
+
+    if not input_otp or not os.path.exists(LOCAL_CSV):
+        return jsonify({"status": "error", "message": "Invalid request"}), 400
+
+    df = pd.read_csv(LOCAL_CSV)
 
-        with sqlite3.connect(DATABASE) as conn:
-            c = conn.cursor()
-            c.execute("INSERT INTO otps (otp, generated_at, expires_at, ip_address, status) VALUES (?, ?, ?, ?, ?)",
-                      (otp, now.isoformat(), expires.isoformat(), ip, "generated"))
-            conn.commit()
+    match = df[(df["otp"] == input_otp) & (df["used"] == False)]
+    if match.empty:
+        return jsonify({"status": "error", "message": "Invalid or used OTP"}), 403
 
-        otp_row = [otp, now.isoformat(), expires.isoformat(), None, ip, "generated"]
-        upload_to_hf(otp_row)
-        flash(f"✅ OTP Generated: {otp}", "success")
+    expiry_time = datetime.fromisoformat(match.iloc[0]["expiry"])
+    if datetime.utcnow() > expiry_time:
+        return jsonify({"status": "error", "message": "OTP expired"}), 403
 
-    return render_template("dashboard.html")
+    df.loc[df["otp"] == input_otp, "used"] = True
+    df.to_csv(LOCAL_CSV, index=False)
+
+    try:
+        dataset = HFDataset.from_pandas(df)
+        dataset.push_to_hub(REPO_ID, token=HF_TOKEN)
+    except Exception as e:
+        print("⚠️ Could not update dataset:", e)
+
+    return jsonify({"status": "success", "message": "OTP verified"})
 
 @app.route("/logout")
 def logout():
     session.clear()
     return redirect(url_for("index"))
 
+# --- Prevent iframe login issues ---
+@app.after_request
+def add_header(response):
+    response.headers["X-Frame-Options"] = "SAMEORIGIN"
+    return response
+
+# --- Start ---
 if __name__ == "__main__":
-    init_db()
     app.run(host="0.0.0.0", port=7860)