MailPocket/shared/admin.go

package shared

import (
	"database/sql"
	"encoding/json"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"

	_ "modernc.org/sqlite"
)

var (
	adminOnce   sync.Once
	adminServer *http.Server
)

type AdminServer struct {
	DB   *sql.DB
	auth *AdminAuth
}

type AdminAuth struct {
	Password string
}

func (a *AdminAuth) Middleware(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		provided := r.Header.Get("X-Admin-Password")
		if provided != a.Password {
			http.Error(w, "Unauthorized", http.StatusUnauthorized)
			return
		}
		next(w, r)
	}
}

func StartAdminServer(dbdir string) {
	var port, _, _ = GetArgs()
	
	adminOnce.Do(func() {
		adminPassword := os.Getenv("ADMIN_PASSWORD")
		if adminPassword == "" {
			log.Fatal("ADMIN_PASSWORD environment variable required")
		}

		db, err := sql.Open("sqlite", filepath.Join(dbdir, "admin.db"))
		if err != nil {
			log.Fatal("Failed to open admin database:", err)
		}

		createTables(db)

		as := &AdminServer{
			DB:   db,
			auth: &AdminAuth{Password: adminPassword},
		}

		mux := http.NewServeMux()
		mux.HandleFunc("/forms/", as.handleForms)
		mux.HandleFunc("/submit/", as.handleSubmit)

		adminServer = &http.Server{
			Addr:    ":" + port,
			Handler: mux,
		}

		log.Println("Starting admin server on port", port)
		go func() {
			if err := adminServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
				log.Fatal("Admin server failed:", err)
			}
		}()
	})
}

func createTables(db *sql.DB) {
	db.Exec(`CREATE TABLE IF NOT EXISTS forms (
		name TEXT PRIMARY KEY,
		created TIMESTAMP DEFAULT CURRENT_TIMESTAMP
	)`)

	db.Exec(`CREATE TABLE IF NOT EXISTS submissions (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		form_name TEXT,
		data TEXT,
		created TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
		FOREIGN KEY(form_name) REFERENCES forms(name) ON DELETE CASCADE
	)`)
}

func (as *AdminServer) handleForms(w http.ResponseWriter, r *http.Request) {
	formName := strings.TrimPrefix(r.URL.Path, "/forms/")

	switch r.Method {
	case http.MethodPost:
		as.auth.Middleware(func(w http.ResponseWriter, r *http.Request) {
			_, err := as.DB.Exec("INSERT INTO forms(name) VALUES (?)", formName)
			if err != nil {
				http.Error(w, "Form creation failed", http.StatusBadRequest)
				return
			}
			w.WriteHeader(http.StatusCreated)
		})(w, r)

	case http.MethodDelete:
		as.auth.Middleware(func(w http.ResponseWriter, r *http.Request) {
			_, err := as.DB.Exec("DELETE FROM forms WHERE name = ?", formName)
			if err != nil {
				http.Error(w, "Form deletion failed", http.StatusInternalServerError)
				return
			}
			w.WriteHeader(http.StatusNoContent)
		})(w, r)

	case http.MethodGet:
		var exists bool
		err := as.DB.QueryRow("SELECT EXISTS(SELECT 1 FROM forms WHERE name = ?)", formName).Scan(&exists)
		if err != nil || !exists {
			http.Error(w, "Form not found", http.StatusNotFound)
			return
		}

		rows, err := as.DB.Query("SELECT data FROM submissions WHERE form_name = ?", formName)
		if err != nil {
			http.Error(w, "Failed to retrieve submissions", http.StatusInternalServerError)
			return
		}
		defer rows.Close()

		var submissions []string
		for rows.Next() {
			var data string
			if err := rows.Scan(&data); err != nil {
				continue
			}
			submissions = append(submissions, data)
		}
		json.NewEncoder(w).Encode(submissions)

	default:
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
	}
}

func (as *AdminServer) handleSubmit(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	formName := strings.TrimPrefix(r.URL.Path, "/submit/")
	data := r.FormValue("data")

	var exists bool
	err := as.DB.QueryRow("SELECT EXISTS(SELECT 1 FROM forms WHERE name = ?)", formName).Scan(&exists)
	if err != nil || !exists {
		http.Error(w, "Form does not exist", http.StatusNotFound)
		return
	}

	_, err = as.DB.Exec(`INSERT INTO submissions(form_name, data) VALUES (?, ?)`, formName, data)
	if err != nil {
		http.Error(w, "Submission failed", http.StatusInternalServerError)
		return
	}

	w.WriteHeader(http.StatusCreated)
}