wg/internal/handlers/handlers.go

package handlers

import (
	"archive/zip"
	"bytes"
	"context"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"html/template"
	"io"
	"log"
	"math/big"
	"net/http"
	"net/url"
	"path"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/evilfox/wg-panel/internal/auth"
	"github.com/evilfox/wg-panel/internal/config"
	"github.com/evilfox/wg-panel/internal/models"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/skip2/go-qrcode"
)

type App struct {
	DB        *pgxpool.Pool
	Auth      *auth.Store
	Cfg       config.Config
	Templates *template.Template
}

func New(db *pgxpool.Pool, a *auth.Store, cfg config.Config, templatesDir string) (*App, error) {
	funcs := template.FuncMap{
		"appURL":  func() string { return cfg.AppURL },
		"siteName": func() string { return cfg.SiteName },
		"add":     func(a, b int) int { return a + b },
		"sub":     func(a, b int) int { return a - b },
	}
	tmpl, err := template.New("").Funcs(funcs).ParseGlob(filepath.Join(templatesDir, "*.html"))
	if err != nil {
		return nil, err
	}
	return &App{DB: db, Auth: a, Cfg: cfg, Templates: tmpl}, nil
}

func (a *App) render(w http.ResponseWriter, name string, data any) {
	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	if err := a.Templates.ExecuteTemplate(w, name, data); err != nil {
		log.Printf("template %s: %v", name, err)
		http.Error(w, "template error", http.StatusInternalServerError)
	}
}

func (a *App) SiteOnline(ctx context.Context) (models.SiteStatus, error) {
	var s models.SiteStatus
	err := a.DB.QueryRow(ctx, `SELECT id, is_online, message FROM site_status WHERE id = 1`).Scan(&s.ID, &s.IsOnline, &s.Message)
	if err != nil {
		return models.SiteStatus{IsOnline: true, Message: ""}, err
	}
	return s, nil
}

func (a *App) OfflineMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if strings.HasPrefix(r.URL.Path, "/login") || strings.HasPrefix(r.URL.Path, "/admin") || strings.HasPrefix(r.URL.Path, "/static") {
			next.ServeHTTP(w, r)
			return
		}
		u := a.Auth.CurrentUser(r)
		if u != nil && u.Role == "admin" {
			next.ServeHTTP(w, r)
			return
		}
		st, err := a.SiteOnline(r.Context())
		if err == nil && !st.IsOnline {
			a.render(w, "offline.html", map[string]any{"Message": st.Message, "SiteName": a.Cfg.SiteName})
			return
		}
		next.ServeHTTP(w, r)
	})
}

func (a *App) Home(w http.ResponseWriter, r *http.Request) {
	u := a.Auth.CurrentUser(r)
	a.render(w, "home.html", map[string]any{"User": u})
}

func (a *App) LoginPage(w http.ResponseWriter, r *http.Request) {
	if u := a.Auth.CurrentUser(r); u != nil {
		if u.Role == "admin" {
			http.Redirect(w, r, "/admin", http.StatusSeeOther)
			return
		}
		http.Redirect(w, r, "/", http.StatusSeeOther)
		return
	}
	a.render(w, "login.html", map[string]any{"Error": ""})
}

func (a *App) LoginPost(w http.ResponseWriter, r *http.Request) {
	if err := r.ParseForm(); err != nil {
		a.render(w, "login.html", map[string]any{"Error": "Некорректная форма"})
		return
	}
	username := strings.TrimSpace(r.FormValue("username"))
	password := r.FormValue("password")
	remember := r.FormValue("remember_me") == "1"
	u, err := a.Auth.Authenticate(r.Context(), username, password)
	if err != nil {
		msg := "Неверное имя пользователя или пароль."
		if err.Error() == "account pending approval" {
			msg = "Ваш аккаунт ожидает одобрения администратором."
		}
		a.render(w, "login.html", map[string]any{"Error": msg, "Username": username})
		return
	}
	if err := a.Auth.Login(w, r, u, remember); err != nil {
		a.render(w, "login.html", map[string]any{"Error": "Ошибка сессии"})
		return
	}
	if u.Role == "admin" {
		http.Redirect(w, r, "/admin", http.StatusSeeOther)
		return
	}
	http.Redirect(w, r, "/", http.StatusSeeOther)
}

func (a *App) Logout(w http.ResponseWriter, r *http.Request) {
	_ = a.Auth.Logout(w, r)
	http.Redirect(w, r, "/login", http.StatusSeeOther)
}

func (a *App) AdminDashboard(w http.ResponseWriter, r *http.Request) {
	u := a.Auth.CurrentUser(r)
	var totalConfigs, totalUsers int
	_ = a.DB.QueryRow(r.Context(), `SELECT COUNT(*) FROM wg_configs`).Scan(&totalConfigs)
	_ = a.DB.QueryRow(r.Context(), `SELECT COUNT(*) FROM users`).Scan(&totalUsers)
	st, _ := a.SiteOnline(r.Context())
	a.render(w, "admin_dashboard.html", map[string]any{
		"User":         u,
		"TotalConfigs": totalConfigs,
		"TotalUsers":   totalUsers,
		"SiteOnline":   st.IsOnline,
		"SiteMessage":  st.Message,
	})
}

func (a *App) AdminStatusPost(w http.ResponseWriter, r *http.Request) {
	_ = r.ParseForm()
	online := r.FormValue("is_online") == "1"
	message := strings.TrimSpace(r.FormValue("message"))
	if message == "" {
		message = "Сайт временно недоступен."
	}
	_, _ = a.DB.Exec(r.Context(), `
		UPDATE site_status SET is_online = $1, message = $2, updated_at = NOW() WHERE id = 1`, online, message)
	http.Redirect(w, r, "/admin", http.StatusSeeOther)
}

func generateToken(n int) (string, error) {
	const chars = "abcdefghijkmnopqrstuvwxyzABCDEFGHJKLMNPQRSTUVWXYZ23456789"
	b := make([]byte, n)
	for i := 0; i < n; i++ {
		v, err := rand.Int(rand.Reader, big.NewInt(int64(len(chars))))
		if err != nil {
			return "", err
		}
		b[i] = chars[v.Int64()]
	}
	return string(b), nil
}

func (a *App) AdminWG(w http.ResponseWriter, r *http.Request) {
	u := a.Auth.CurrentUser(r)
	page, _ := strconv.Atoi(r.URL.Query().Get("page"))
	if page < 1 {
		page = 1
	}
	const perPage = 10
	var total int
	_ = a.DB.QueryRow(r.Context(), `SELECT COUNT(*) FROM wg_configs`).Scan(&total)
	totalPages := (total + perPage - 1) / perPage
	if totalPages < 1 {
		totalPages = 1
	}
	offset := (page - 1) * perPage

	rows, err := a.DB.Query(r.Context(), `
		SELECT id, COALESCE(token,''), original_filename, created_at
		FROM wg_configs ORDER BY created_at DESC LIMIT $1 OFFSET $2`, perPage, offset)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	defer rows.Close()
	var configs []models.WGConfig
	for rows.Next() {
		var c models.WGConfig
		if err := rows.Scan(&c.ID, &c.Token, &c.OriginalFilename, &c.CreatedAt); err != nil {
			http.Error(w, err.Error(), 500)
			return
		}
		configs = append(configs, c)
	}

	msg := r.URL.Query().Get("msg")
	msgType := r.URL.Query().Get("type")
	token := r.URL.Query().Get("token")
	settings := a.wgSettings(r.Context())
	a.render(w, "admin_wg.html", map[string]any{
		"User":        u,
		"Configs":     configs,
		"Page":        page,
		"TotalPages":  totalPages,
		"Message":     msg,
		"MessageType": msgType,
		"NewToken":    token,
		"PublicURL":   a.Cfg.AppURL + "/wg?token=",
		"ShowZipFull": settings["show_zip_full"] == "1",
		"ShowZipConf": settings["show_zip_conf"] == "1",
	})
}

func (a *App) AdminWGUpload(w http.ResponseWriter, r *http.Request) {
	if err := r.ParseMultipartForm(1 << 20); err != nil {
		http.Redirect(w, r, "/admin/wg?msg="+urlQ("Ошибка формы")+"&type=error", http.StatusSeeOther)
		return
	}
	file, hdr, err := r.FormFile("wg_conf")
	if err != nil {
		http.Redirect(w, r, "/admin/wg?msg="+urlQ("Файл не выбран")+"&type=error", http.StatusSeeOther)
		return
	}
	defer file.Close()
	if !strings.HasSuffix(strings.ToLower(hdr.Filename), ".conf") {
		http.Redirect(w, r, "/admin/wg?msg="+urlQ("Разрешены только .conf")+"&type=error", http.StatusSeeOther)
		return
	}
	content, err := io.ReadAll(io.LimitReader(file, 10000))
	if err != nil || len(content) == 0 {
		http.Redirect(w, r, "/admin/wg?msg="+urlQ("Не удалось прочитать файл")+"&type=error", http.StatusSeeOther)
		return
	}

	var token string
	for {
		token, err = generateToken(5)
		if err != nil {
			http.Redirect(w, r, "/admin/wg?msg="+urlQ("Ошибка токена")+"&type=error", http.StatusSeeOther)
			return
		}
		var exists bool
		_ = a.DB.QueryRow(r.Context(), `SELECT EXISTS(SELECT 1 FROM wg_configs WHERE token=$1)`, token).Scan(&exists)
		if !exists {
			break
		}
	}
	orig := path.Base(hdr.Filename)
	_, err = a.DB.Exec(r.Context(), `
		INSERT INTO wg_configs (token, config_content, original_filename)
		VALUES ($1, $2, $3)`, token, string(content), orig)
	if err != nil {
		http.Redirect(w, r, "/admin/wg?msg="+urlQ("Ошибка сохранения")+"&type=error", http.StatusSeeOther)
		return
	}
	http.Redirect(w, r, "/admin/wg?msg="+urlQ("Конфиг сохранён")+"&type=success&token="+token, http.StatusSeeOther)
}

func (a *App) AdminWGDelete(w http.ResponseWriter, r *http.Request) {
	_ = r.ParseForm()
	id, _ := strconv.ParseInt(r.FormValue("id"), 10, 64)
	if id > 0 {
		_, _ = a.DB.Exec(r.Context(), `DELETE FROM wg_configs WHERE id=$1`, id)
	}
	http.Redirect(w, r, "/admin/wg", http.StatusSeeOther)
}

func (a *App) PublicWG(w http.ResponseWriter, r *http.Request) {
	token := sanitizeToken(r.URL.Query().Get("token"))
	if len(token) < 5 {
		http.Error(w, "Неверный токен", http.StatusBadRequest)
		return
	}
	download := r.URL.Query().Get("download")

	var cfg models.WGConfig
	err := a.DB.QueryRow(r.Context(), `
		SELECT id, COALESCE(token,''), config_content, original_filename
		FROM wg_configs WHERE token=$1`, token,
	).Scan(&cfg.ID, &cfg.Token, &cfg.ConfigContent, &cfg.OriginalFilename)
	if err != nil {
		if err == pgx.ErrNoRows {
			http.Error(w, "Конфиг не найден", http.StatusNotFound)
			return
		}
		http.Error(w, err.Error(), 500)
		return
	}
	fileName := sanitizeFilename(cfg.OriginalFilename)
	if !strings.HasSuffix(strings.ToLower(fileName), ".conf") {
		fileName += ".conf"
	}

	settings := a.wgSettings(r.Context())

	if download == "full" || download == "conf" {
		var buf bytes.Buffer
		zw := zip.NewWriter(&buf)
		fw, err := zw.Create(fileName)
		if err != nil {
			http.Error(w, "zip error", 500)
			return
		}
		_, _ = fw.Write([]byte(cfg.ConfigContent))
		if download == "full" {
			png, err := qrcode.Encode(cfg.ConfigContent, qrcode.Medium, 300)
			if err == nil {
				qrw, _ := zw.Create("qr_code.png")
				_, _ = qrw.Write(png)
			}
		}
		_ = zw.Close()
		name := fmt.Sprintf("wg_config_%s.zip", token)
		if download == "conf" {
			name = fmt.Sprintf("wg_config_%s_only.conf.zip", token)
		}
		w.Header().Set("Content-Type", "application/zip")
		w.Header().Set("Content-Disposition", `attachment; filename="`+name+`"`)
		_, _ = w.Write(buf.Bytes())
		return
	}

	png, err := qrcode.Encode(cfg.ConfigContent, qrcode.Medium, 300)
	qrData := ""
	if err == nil {
		qrData = "data:image/png;base64," + base64.StdEncoding.EncodeToString(png)
	}

	a.render(w, "wg_public.html", map[string]any{
		"FileName":    fileName,
		"Token":       token,
		"QR":          qrData,
		"Config":      cfg.ConfigContent,
		"ShowZipFull": settings["show_zip_full"] == "1",
		"ShowZipConf": settings["show_zip_conf"] == "1",
	})
}

func (a *App) wgSettings(ctx context.Context) map[string]string {
	out := map[string]string{"show_zip_full": "0", "show_zip_conf": "1"}
	rows, err := a.DB.Query(ctx, `SELECT setting_key, setting_value FROM wg_config_settings`)
	if err != nil {
		return out
	}
	defer rows.Close()
	for rows.Next() {
		var k, v string
		if rows.Scan(&k, &v) == nil {
			out[k] = v
		}
	}
	return out
}

func (a *App) AdminWGSettings(w http.ResponseWriter, r *http.Request) {
	_ = r.ParseForm()
	full := "0"
	conf := "0"
	if r.FormValue("show_zip_full") == "1" {
		full = "1"
	}
	if r.FormValue("show_zip_conf") == "1" {
		conf = "1"
	}
	_, _ = a.DB.Exec(r.Context(), `
		INSERT INTO wg_config_settings (setting_key, setting_value, updated_at)
		VALUES ('show_zip_full', $1, NOW())
		ON CONFLICT (setting_key) DO UPDATE SET setting_value = EXCLUDED.setting_value, updated_at = NOW()`, full)
	_, _ = a.DB.Exec(r.Context(), `
		INSERT INTO wg_config_settings (setting_key, setting_value, updated_at)
		VALUES ('show_zip_conf', $1, NOW())
		ON CONFLICT (setting_key) DO UPDATE SET setting_value = EXCLUDED.setting_value, updated_at = NOW()`, conf)
	http.Redirect(w, r, "/admin/wg", http.StatusSeeOther)
}

func sanitizeToken(s string) string {
	var b strings.Builder
	for _, r := range s {
		if (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') || (r >= '0' && r <= '9') {
			b.WriteRune(r)
		}
	}
	return b.String()
}

func sanitizeFilename(s string) string {
	s = path.Base(s)
	var b strings.Builder
	for _, r := range s {
		if (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') || (r >= '0' && r <= '9') || r == '.' || r == '_' || r == '-' {
			b.WriteRune(r)
		} else {
			b.WriteByte('_')
		}
	}
	if b.Len() == 0 {
		return "wg_config.conf"
	}
	return b.String()
}

func urlQ(s string) string {
	return url.QueryEscape(s)
}