add password/mfa setup to profile

handlers.go

@@ -1,24 +1,160 @@
 package main
 
 import (
+	"encoding/json"
 	"net/http"
 	"regexp"
 	"sort"
 	"strconv"
 	"strings"
+
+	"golang.org/x/crypto/bcrypt"
 )
 
+const sessionName = "session"
+
 func profileHandler(w http.ResponseWriter, r *http.Request) {
 	configMu.RLock()
 	mfaEnabled := config.MFASecret != ""
 	configMu.RUnlock()
 
-	data := struct {
+	session, _ := store.Get(r, sessionName)
-		MFAEnabled bool
+
-	}{
+	data := TemplateData{
 		MFAEnabled: mfaEnabled,
 	}
+
+	// Check if MFA setup is in progress
+	if setupSecret, ok := session.Values["mfa_setup_secret"].(string); ok && setupSecret != "" {
+		if setupURL, ok := session.Values["mfa_setup_url"].(string); ok {
+			data.MFASetupInProgress = true
+			data.MFASecret = setupSecret
+			data.MFAURL = setupURL
+		}
+	}
+
+	if r.Method == http.MethodPost {
+		action := r.FormValue("action")
+		switch action {
+		case "change_password":
+			current := r.FormValue("current_password")
+			newpw := r.FormValue("new_password")
+			confirmpw := r.FormValue("confirm_new_password")
+			mfaCode := r.FormValue("mfa_code")
+
+			// Validate current password
+			if bcrypt.CompareHashAndPassword([]byte(config.HashedPass), []byte(current)) != nil {
+				data.Notification = "Current password is incorrect."
+				data.NotificationType = "error"
 				renderTemplate(w, r, "profile.html", data)
+				return
+			}
+			// Check if new passwords match
+			if newpw != confirmpw {
+				data.Notification = "New passwords do not match."
+				data.NotificationType = "error"
+				renderTemplate(w, r, "profile.html", data)
+				return
+			}
+			// If MFA enabled, validate code
+			if config.MFASecret != "" {
+				if !validateTOTP(config.MFASecret, mfaCode) {
+					data.Notification = "Invalid MFA code."
+					data.NotificationType = "error"
+					renderTemplate(w, r, "profile.html", data)
+					return
+				}
+			}
+			// Validate new password
+			if len(newpw) < 8 {
+				data.Notification = "Password must be at least 8 characters."
+				data.NotificationType = "error"
+				renderTemplate(w, r, "profile.html", data)
+				return
+			}
+			hash, _ := bcrypt.GenerateFromPassword([]byte(newpw), bcrypt.DefaultCost)
+			config.HashedPass = string(hash)
+			saveConfig()
+			// Invalidate session
+			session.Options.MaxAge = -1
+			session.Save(r, w)
+			data.Notification = "Password changed successfully. Please log in again."
+			data.NotificationType = "success"
+			http.Redirect(w, r, "/login", http.StatusSeeOther)
+			return
+
+		case "toggle_mfa":
+			mfaAction := r.FormValue("mfa")
+			if mfaAction == "on" && config.MFASecret == "" {
+				// Generate new MFA secret and store in session
+				secret, qrURL := generateMFASecret(config.Username, config.AppName)
+				session.Values["mfa_setup_secret"] = secret
+				session.Values["mfa_setup_url"] = qrURL
+				session.Save(r, w)
+				data.MFASetupInProgress = true
+				data.MFASecret = secret
+				data.MFAURL = qrURL
+				data.Notification = "MFA setup initiated. Please configure your authenticator app."
+				data.NotificationType = "info"
+			} else if mfaAction == "off" && config.MFASecret != "" {
+				config.MFASecret = ""
+				saveConfig()
+				data.Notification = "MFA disabled."
+				data.NotificationType = "success"
+				data.MFAEnabled = false
+			}
+
+		case "confirm_mfa":
+			mfaCode := r.FormValue("mfa_code")
+			if setupSecret, ok := session.Values["mfa_setup_secret"].(string); ok && setupSecret != "" {
+				if validateTOTP(setupSecret, mfaCode) {
+					config.MFASecret = setupSecret
+					saveConfig()
+					// Clear session
+					delete(session.Values, "mfa_setup_secret")
+					delete(session.Values, "mfa_setup_url")
+					session.Save(r, w)
+					data.Notification = "MFA enabled successfully."
+					data.NotificationType = "success"
+					data.MFAEnabled = true
+					data.MFASetupInProgress = false
+					data.MFASecret = ""
+					data.MFAURL = ""
+				} else {
+					data.Notification = "Invalid MFA code. Please try again."
+					data.NotificationType = "error"
+				}
+			}
+
+		case "cancel_mfa":
+			// Clear session
+			delete(session.Values, "mfa_setup_secret")
+			delete(session.Values, "mfa_setup_url")
+			session.Save(r, w)
+			data.MFASetupInProgress = false
+			data.MFASecret = ""
+			data.MFAURL = ""
+			data.Notification = "MFA setup cancelled."
+			data.NotificationType = "info"
+		}
+	}
+
+	renderTemplate(w, r, "profile.html", data)
+}
+
+func mfaSetupHandler(w http.ResponseWriter, r *http.Request) {
+	if config.MFASecret != "" {
+		http.Error(w, "MFA already enabled", 400)
+		return
+	}
+	secret, qrURL := generateMFASecret(config.Username, config.AppName)
+	// Store in session
+	session, _ := store.Get(r, sessionName)
+	session.Values["mfa_setup_secret"] = secret
+	session.Values["mfa_setup_url"] = qrURL
+	session.Save(r, w)
+	w.Header().Set("Content-Type", "application/json")
+	json.NewEncoder(w).Encode(map[string]string{"secret": secret, "url": qrURL})
 }
 
 func dashboardHandler(w http.ResponseWriter, r *http.Request) {

main.go

@@ -64,6 +64,7 @@ func main() {
 
 	// Protected routes
 	r.HandleFunc("/profile", authMiddleware(profileHandler)).Methods("GET", "POST")
+	r.HandleFunc("/api/mfa-setup", authMiddleware(mfaSetupHandler)).Methods("POST")
 	r.HandleFunc("/", authMiddleware(dashboardHandler))
 	r.HandleFunc("/urllists", authMiddleware(urlListsHandler)).Methods("GET", "POST")
 	r.HandleFunc("/domains", authMiddleware(domainsHandler)).Methods("GET", "POST")

templates/profile.html

@@ -1,22 +1,93 @@
 {{ define "content" }}
 <div class="max-w-2xl mx-auto">
     <h1 class="text-3xl font-bold mb-6">Profile</h1>
-    <div class="bg-gray-800 p-6 rounded">
+    <div class="bg-gray-800 p-6 rounded mb-6">
         <p class="mb-4">
             <strong class="text-white">Username:</strong> 
             <span class="text-gray-300">admin</span>
         </p>
         <p class="mb-6">
             <strong class="text-white">MFA Status:</strong> 
-            {{ if .PageData.MFAEnabled }}
+            {{ if .MFAEnabled }}
                 <span class="text-green-400 font-bold">Enabled</span>
             {{ else }}
                 <span class="text-yellow-400 font-bold">Disabled</span>
             {{ end }}
         </p>
-        <div class="bg-gray-700 p-4 rounded mt-4">
+    </div>
-            <p class="text-sm text-gray-300">
+
-                <strong>To manage MFA or change password:</strong>
+    <!-- Password Change Form -->
+    <div class="bg-gray-800 p-6 rounded mb-6">
+        <h2 class="text-xl font-bold mb-4 text-white">Change Password</h2>
+        <form method="POST" action="/profile">
+            <input type="hidden" name="action" value="change_password">
+            <input type="hidden" name="csrf_token" value="{{ .CSRFToken }}">
+            <div class="mb-4">
+                <label class="block text-gray-300 mb-1">Current Password</label>
+                <input type="password" name="current_password" required class="w-full p-2 rounded bg-gray-700 text-white">
+            </div>
+            <div class="mb-4">
+                <label class="block text-gray-300 mb-1">New Password</label>
+                <input type="password" name="new_password" required class="w-full p-2 rounded bg-gray-700 text-white">
+            </div>
+            <div class="mb-4">
+                <label class="block text-gray-300 mb-1">Confirm New Password</label>
+                <input type="password" name="confirm_new_password" required class="w-full p-2 rounded bg-gray-700 text-white">
+            </div>
+            {{ if .MFAEnabled }}
+            <div class="mb-4">
+                <label class="block text-gray-300 mb-1">MFA Code</label>
+                <input type="text" name="mfa_code" required class="w-full p-2 rounded bg-gray-700 text-white">
+            </div>
+            {{ end }}
+            <button type="submit" class="bg-blue-600 hover:bg-blue-700 text-white px-4 py-2 rounded">Change Password</button>
+        </form>
+    </div>
+
+    <!-- MFA Toggle Form -->
+    <div class="bg-gray-800 p-6 rounded mb-6" id="mfa-toggle">
+        <h2 class="text-xl font-bold mb-4 text-white">MFA Management</h2>
+        <form method="POST" action="/profile">
+            <input type="hidden" name="action" value="toggle_mfa">
+            <input type="hidden" name="csrf_token" value="{{ .CSRFToken }}">
+            {{ if .MFAEnabled }}
+                <button type="submit" name="mfa" value="off" class="bg-red-600 hover:bg-red-700 text-white px-4 py-2 rounded">Disable MFA</button>
+            {{ else if not .MFASetupInProgress }}
+                <button type="button" onclick="enableMFA()" class="bg-green-600 hover:bg-green-700 text-white px-4 py-2 rounded">Enable MFA</button>
+            {{ end }}
+        </form>
+    </div>
+
+    <!-- MFA Setup Section -->
+    <div class="bg-gray-800 p-6 rounded mb-6" id="mfa-setup" style="display: {{ if .MFASetupInProgress }}block{{ else }}none{{ end }};">
+        <h2 class="text-xl font-bold mb-4 text-white">Setup MFA</h2>
+        <p class="text-gray-300 mb-4">Scan the QR code with your authenticator app or enter the secret manually.</p>
+        <div class="mb-4">
+            <strong class="text-white">Secret:</strong> <code class="bg-gray-900 px-2 py-1 rounded" id="mfa-secret">{{ .MFASecret }}</code>
+        </div>
+        <div class="mb-4">
+            <strong class="text-white">OTP URL:</strong> <code class="bg-gray-900 px-2 py-1 rounded" id="mfa-url">{{ .MFAURL }}</code>
+        </div>
+        <div class="mb-4">
+            <img id="mfa-qr" src="https://api.qrserver.com/v1/create-qr-code/?size=150x150&data={{ .MFAURL }}" alt="QR Code" class="border border-gray-600" style="display: {{ if .MFASetupInProgress }}block{{ else }}none{{ end }};">
+        </div>
+        <form method="POST" action="/profile">
+            <input type="hidden" name="action" value="confirm_mfa">
+            <input type="hidden" name="csrf_token" value="{{ .CSRFToken }}">
+            <div class="mb-4">
+                <label class="block text-gray-300 mb-1">Enter 6-digit code from your app</label>
+                <input type="text" name="mfa_code" required class="w-full p-2 rounded bg-gray-700 text-white">
+            </div>
+            <button type="submit" class="bg-green-600 hover:bg-green-700 text-white px-4 py-2 rounded mr-2">Confirm MFA</button>
+            <button type="submit" name="action" value="cancel_mfa" class="bg-red-600 hover:bg-red-700 text-white px-4 py-2 rounded">Cancel</button>
+        </form>
+    </div>
+
+    <!-- Terminal Commands -->
+    <div class="bg-gray-800 p-6 rounded">
+        <h2 class="text-xl font-bold mb-4 text-white">Terminal Commands</h2>
+        <p class="text-sm text-gray-300 mb-2">
+            <strong>Alternative CLI management:</strong>
         </p>
         <p class="text-sm text-gray-400 mt-2">Use the command-line interface with the executable flags:</p>
         <ul class="text-sm text-gray-400 mt-2 ml-4 space-y-1">
@@ -25,6 +96,27 @@
             <li>• <code class="bg-gray-900 px-2 py-1 rounded">./unifi-blocklist-app -mfa off</code> - Disable MFA</li>
         </ul>
     </div>
-    </div>
 </div>
+<script>
+function enableMFA() {
+    fetch('/api/mfa-setup', {
+        method: 'POST',
+        headers: {
+            'X-CSRF-Token': '{{ .CSRFToken }}'
+        }
+    })
+    .then(response => response.json())
+    .then(data => {
+        document.getElementById('mfa-secret').textContent = data.secret;
+        document.getElementById('mfa-url').textContent = data.url;
+        document.getElementById('mfa-qr').src = 'https://api.qrserver.com/v1/create-qr-code/?size=150x150&data=' + encodeURIComponent(data.url);
+        document.getElementById('mfa-qr').style.display = 'block';
+        document.getElementById('mfa-setup').style.display = 'block';
+        document.getElementById('mfa-toggle').style.display = 'none';
+    })
+    .catch(error => {
+        alert('Error enabling MFA: ' + error);
+    });
+}
+</script>
 {{ end }}

utils.go

@@ -12,6 +12,7 @@ import (
 	"strings"
 
 	"github.com/justinas/nosurf"
+	"github.com/pquerna/otp/totp"
 )
 
 //go:embed templates/*.html static/*
@@ -24,6 +25,10 @@ type TemplateData struct {
 	AppName             string
 	Notification        string
 	NotificationType    string
+	MFAEnabled          bool
+	MFASetupInProgress  bool
+	MFASecret           string
+	MFAURL              string
 }
 
 func renderTemplate(w http.ResponseWriter, r *http.Request, tmpl string, pageData interface{}) {
@@ -36,6 +41,7 @@ func renderTemplate(w http.ResponseWriter, r *http.Request, tmpl string, pageDat
 
 	configMu.RLock()
 	appName := config.AppName
+	mfaEnabled := config.MFASecret != ""
 	configMu.RUnlock()
 
 	td := TemplateData{
@@ -43,6 +49,7 @@ func renderTemplate(w http.ResponseWriter, r *http.Request, tmpl string, pageDat
 		PageData:      pageData,
 		Authenticated: isAuth,
 		AppName:       appName,
+		MFAEnabled:    mfaEnabled,
 	}
 
 	// Extract notification and type from pageData if it has those fields
@@ -249,3 +256,17 @@ func generateSecret() string {
 	// since we're using pquerna/otp for TOTP generation
 	return ""
 }
+
+// validateTOTP validates a TOTP code against a secret
+func validateTOTP(secret, code string) bool {
+	return totp.Validate(code, secret)
+}
+
+// generateMFASecret generates a new MFA secret and QR code URL
+func generateMFASecret(username, appname string) (string, string) {
+	key, _ := totp.Generate(totp.GenerateOpts{
+		Issuer:      appname,
+		AccountName: username,
+	})
+	return key.Secret(), key.URL()
+}