Update server.py

server.py

@@ -1,175 +1,159 @@
 import os
 import sqlite3
-import csv
-import json
+import secrets
 import pathlib
-import random
 from datetime import datetime, timedelta
-from flask import Flask, request, render_template_string, redirect, url_for, session, send_file, jsonify
+from flask import Flask, request, redirect, url_for, render_template, session, flash
 from flask_limiter import Limiter
 from flask_limiter.util import get_remote_address
-from datasets import Dataset
 from huggingface_hub import HfApi
+import pandas as pd
 
-# === Config ===
+# ENV values
 ADMIN_USER = os.getenv("ADMIN_USER", "admin")
 ADMIN_PASS = os.getenv("ADMIN_PASS", "Welcome123")
-FLASK_SECRET = os.getenv("FLASK_SECRET", "IDONTKNOWWHATISTHIS")
-HF_TOKEN = os.getenv("HF_TOKEN", "")
-HF_DATASET_REPO = "0vergeared/otp-logs"
+FLASK_SECRET = os.getenv("FLASK_SECRET", "supersecret")
+HF_TOKEN = os.getenv("HF_TOKEN")
+HF_REPO_ID = "0vergeared/otp-logs"
 
-DATABASE = "/tmp/otp.db"
-EXPORT_CSV = "/tmp/otp_logs.csv"
-EXPORT_JSON = "/tmp/otp_logs.json"
-pathlib.Path("/tmp").mkdir(parents=True, exist_ok=True)
+# Paths
+DATABASE = "otp.db"
+TEMPLATES = "templates"
+STATIC = "static"
 
-# === App Setup ===
-app = Flask(__name__)
+# Setup
+app = Flask(__name__, template_folder=TEMPLATES, static_folder=STATIC)
 app.secret_key = FLASK_SECRET
 
-limiter = Limiter(key_func=get_remote_address)
-limiter.init_app(app)
+# Rate Limiting
+limiter = Limiter(get_remote_address, app=app)
 
-# === Helpers ===
-def generate_otp():
-    return str(random.randint(100000, 999999))
+# Make sure Hugging Face Hub doesn't cache
+os.environ["HF_DATASETS_CACHE"] = "/tmp/hf_cache"
+
+# Make sure templates/static folders exist
+os.makedirs(TEMPLATES, exist_ok=True)
+os.makedirs(STATIC, exist_ok=True)
 
 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,
-                generated_at TEXT,
-                expires_at TEXT,
-                used_at TEXT,
-                ip_address TEXT,
-                status TEXT
-            )
-        """)
+        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()
 
-def export_logs():
-    with sqlite3.connect(DATABASE) as conn:
-        c = conn.cursor()
-        c.execute("SELECT otp, generated_at, expires_at, used_at, ip_address, status FROM otps ORDER BY id DESC")
-        rows = c.fetchall()
-
-        # Save CSV
-        with open(EXPORT_CSV, "w", newline="") as f:
-            writer = csv.writer(f)
-            writer.writerow(["otp", "generated_at", "expires_at", "used_at", "ip_address", "status"])
-            writer.writerows(rows)
-
-        # Save JSON
-        json_data = [
-            {
-                "otp": row[0],
-                "generated_at": row[1],
-                "expires_at": row[2],
-                "used_at": row[3],
-                "ip_address": row[4],
-                "status": row[5]
-            }
-            for row in rows
-        ]
-        with open(EXPORT_JSON, "w") as f:
-            json.dump(json_data, f, indent=2)
+def generate_otp():
+    return ''.join(secrets.choice("0123456789") for _ in range(6))
 
 def upload_to_hf(otp_row):
+    csv_path = "/tmp/otp_all.csv"
+    file_exists = os.path.exists(csv_path)
+
+    df_new = pd.DataFrame([otp_row], columns=["otp", "generated_at", "expires_at", "used_at", "ip", "status"])
+    if file_exists:
+        df_old = pd.read_csv(csv_path)
+        df = pd.concat([df_old, df_new], ignore_index=True)
+    else:
+        df = df_new
+
+    df.to_csv(csv_path, index=False)
+
     try:
-        import pandas as pd
-        df = pd.DataFrame([otp_row], columns=["otp", "generated_at", "expires_at", "used_at", "ip_address", "status"])
-        dataset = Dataset.from_pandas(df)
-        dataset.push_to_hub(repo_id=HF_DATASET_REPO, token=HF_TOKEN)
-        print("✅ OTP log pushed to Hugging Face")
+        api = HfApi()
+        api.upload_file(
+            path_or_fileobj=csv_path,
+            path_in_repo="otp_all.csv",
+            repo_id=HF_REPO_ID,
+            token=HF_TOKEN,
+            create_pr=False
+        )
+        print("✅ OTP log uploaded to HF.")
     except Exception as e:
-        print("⚠️ Failed to upload OTP log:", e)
+        print("⚠️ Failed to upload to HF:", e)
 
-# === Routes ===
-
-@app.route("/")
-def home():
-    if "admin" not in session:
-        return redirect(url_for("login"))
-    return redirect(url_for("generate"))
+@app.route("/", methods=["GET", "POST"])
+def index():
+    if request.method == "POST":
+        otp_input = request.form.get("otp")
+        now = datetime.utcnow()
 
-@app.route("/login", methods=["GET", "POST"])
-def login():
+        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()
+
+                    # Log to Hugging Face
+                    otp_row = [otp_input, gen_at, exp_at, used_time, request.remote_addr, "used"]
+                    upload_to_hf(otp_row)
+
+                    flash("OTP accepted. Installation can continue!", "success")
+                    return redirect(url_for("index"))
+            else:
+                flash("Invalid OTP!", "danger")
+
+    return render_template("index.html")
+
+@app.route("/admin", methods=["GET", "POST"])
+def admin():
     if request.method == "POST":
-        if request.form["username"] == ADMIN_USER and request.form["password"] == ADMIN_PASS:
+        username = request.form.get("username")
+        password = request.form.get("password")
+        if username == ADMIN_USER and password == ADMIN_PASS:
             session["admin"] = True
-            return redirect(url_for("generate"))
-        return "Login failed"
-    return render_template_string("""
-        <h2>Admin Login</h2>
-        <form method="post">
-            Username: <input name="username"><br>
-            Password: <input type="password" name="password"><br>
-            <input type="submit" value="Login">
-        </form>
-    """)
-
-@app.route("/generate", methods=["GET", "POST"])
-def generate():
-    if "admin" not in session:
-        return redirect(url_for("login"))
+            return redirect(url_for("dashboard"))
+        else:
+            flash("Login failed", "danger")
+    return render_template("admin_login.html")
+
+@app.route("/dashboard", methods=["GET", "POST"])
+def dashboard():
+    if not session.get("admin"):
+        return redirect(url_for("admin"))
 
     if request.method == "POST":
         otp = generate_otp()
         now = datetime.utcnow()
-        expires_at = now + timedelta(minutes=10)
+        expires = now + timedelta(minutes=10)
+        ip = request.remote_addr
 
         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_at.isoformat(), request.remote_addr, "generated"))
+                      (otp, now.isoformat(), expires.isoformat(), ip, "generated"))
             conn.commit()
 
-        # Export logs
-        export_logs()
-
-        # Push to Hugging Face Dataset
-        otp_row = [otp, now.isoformat(), expires_at.isoformat(), None, request.remote_addr, "generated"]
+        # Log to HF
+        otp_row = [otp, now.isoformat(), expires.isoformat(), None, ip, "generated"]
         upload_to_hf(otp_row)
 
-        return f"Generated OTP: <b>{otp}</b> (expires in 10 mins)<br><a href='/generate'>Back</a>"
+        flash(f"✅ OTP Generated: {otp}", "success")
 
-    return render_template_string("""
-        <h2>Generate OTP</h2>
-        <form method="post">
-            <input type="submit" value="Generate OTP">
-        </form>
-        <a href='/download-csv'>Download CSV</a> |
-        <a href='/download-json'>Download JSON</a> |
-        <a href='/logout'>Logout</a>
-    """)
+    return render_template("dashboard.html")
 
 @app.route("/logout")
 def logout():
-    session.pop("admin", None)
-    return redirect(url_for("login"))
-
-@app.route("/download-csv")
-def download_csv():
-    export_logs()
-    return send_file(EXPORT_CSV, as_attachment=True)
-
-@app.route("/download-json")
-def download_json():
-    export_logs()
-    return send_file(EXPORT_JSON, as_attachment=True)
-
-@app.route("/debug")
-def debug():
-    return jsonify({
-        "db_exists": os.path.exists(DATABASE),
-        "csv_exists": os.path.exists(EXPORT_CSV),
-        "json_exists": os.path.exists(EXPORT_JSON),
-    })
-
-# === Run App ===
+    session.clear()
+    return redirect(url_for("index"))
+
 if __name__ == "__main__":
     init_db()
     app.run(host="0.0.0.0", port=7860)