drive/backend/main.go

package main

import (
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"log"
	"os"
	"time"

	"git.elijahkuntz.com/Elijah/drive/config"
	"git.elijahkuntz.com/Elijah/drive/database"
	"git.elijahkuntz.com/Elijah/drive/handlers"
	"git.elijahkuntz.com/Elijah/drive/middleware"
	"git.elijahkuntz.com/Elijah/drive/workers"
	"github.com/gofiber/fiber/v2"
	"github.com/gofiber/fiber/v2/middleware/cors"
	"github.com/gofiber/fiber/v2/middleware/helmet"
	"github.com/gofiber/fiber/v2/middleware/limiter"
	"github.com/gofiber/fiber/v2/middleware/logger"
	"github.com/gofiber/fiber/v2/middleware/recover"
)

func main() {
	cfg := config.Load()

	// Auto-generate JWT secret if not set
	if cfg.JWTSecret == "" {
		secret, err := generateSecret(32)
		if err != nil {
			log.Fatal("Failed to generate JWT secret:", err)
		}
		cfg.JWTSecret = secret
		log.Println("Warning: JWT_SECRET not set, generated a random one. Set JWT_SECRET env var for persistence across restarts.")
	}

	// Ensure required directories exist
	dirs := []string{
		cfg.DataDir,
		cfg.StorageDir,
		cfg.ThumbnailDir,
		cfg.TrashDir,
		cfg.VersionsDir,
		cfg.BackupDir,
	}
	for _, dir := range dirs {
		if err := os.MkdirAll(dir, 0755); err != nil {
			log.Fatalf("Failed to create directory %s: %v", dir, err)
		}
	}

	// Initialize database
	db, err := database.New(cfg.DBPath)
	if err != nil {
		log.Fatal("Failed to initialize database:", err)
	}
	defer db.Close()

	// Initialize brute force guard
	guard := middleware.NewBruteForceGuard(cfg.MaxLoginAttempts, cfg.LockoutSeconds)

	// Initialize handlers
	authHandler := &handlers.AuthHandler{DB: db, Config: cfg, Guard: guard}
	fsHandler := &handlers.FSHandler{DB: db, Config: cfg}
	healthHandler := &handlers.HealthHandler{DB: db, Config: cfg}
	settingsHandler := &handlers.SettingsHandler{DB: db}
	tusHandler := handlers.NewTusHandler(db, cfg)
	webdavHandler := handlers.NewWebDAVHandler(db, cfg)
	archiveHandler := &handlers.ArchiveHandler{DB: db, Config: cfg}
	shareHandler := &handlers.ShareHandler{DB: db, Config: cfg}

	// Initialize background workers
	workers.InitTaskManager(cfg)
	workers.InitThumbnailManager(cfg, db)
	workers.StartBackupWorker(cfg, db)
	workers.StartTrashWorker(cfg, db)

	// Create Fiber app
	app := fiber.New(fiber.Config{
		BodyLimit:             100 * 1024 * 1024 * 1024, // 100 GB (for chunked uploads, each chunk is smaller)
		StreamRequestBody:    true,
		DisableStartupMessage: false,
		AppName:              "Drive",
	})

	// Global middleware
	app.Use(recover.New())
	app.Use(logger.New(logger.Config{
		Format:     "${time} | ${status} | ${latency} | ${ip} | ${method} | ${path}\n",
		TimeFormat: "2006-01-02 15:04:05",
	}))
	app.Use(func(c *fiber.Ctx) error {
		c.Set("Cache-Control", "no-store, no-cache, must-revalidate")
		return c.Next()
	})
	app.Use(helmet.New())
	app.Use(cors.New(cors.Config{
		AllowOrigins: "*",
		AllowHeaders: "Origin, Content-Type, Accept, Authorization, Upload-Length, Upload-Offset, Upload-Metadata, Tus-Resumable",
		AllowMethods: "GET, POST, PUT, DELETE, PATCH, OPTIONS, HEAD",
		ExposeHeaders: "Location, Upload-Offset, Upload-Length, Tus-Resumable, Tus-Version, Tus-Extension, Tus-Max-Size",
	}))

	// === PUBLIC ROUTES (no auth) ===

	// Health check
	app.Get("/health", healthHandler.Health)

	// Setup routes (first-run only)
	app.Get("/api/setup/status", authHandler.CheckSetup)
	app.Post("/api/setup", authHandler.Setup)

	// Public Share Endpoints (Rate Limited)
	public := app.Group("/api/public")
	public.Use(limiter.New(limiter.Config{
		Max:        20,
		Expiration: 1 * time.Minute,
	}))
	public.Post("/share/:id", shareHandler.GetPublicShare)
	public.Get("/download/:id", shareHandler.DownloadPublicShare)

	// Auth routes (with brute force protection)
	auth := app.Group("/api/auth")
	auth.Post("/login", guard.Check(), authHandler.Login)

	// === PROTECTED ROUTES (require JWT) ===
	protected := app.Group("/api", middleware.AuthMiddleware(cfg.JWTSecret))

	// Token refresh
	protected.Post("/auth/refresh", authHandler.Refresh)

	// Password management
	protected.Put("/auth/password", authHandler.ChangePassword)

	// 2FA management
	protected.Post("/auth/2fa/enable", authHandler.Enable2FA)
	protected.Post("/auth/2fa/confirm", authHandler.Confirm2FA)
	protected.Post("/auth/2fa/disable", authHandler.Disable2FA)

	// Settings & Audit
	protected.Get("/settings", settingsHandler.GetSettings)
	protected.Put("/settings", settingsHandler.UpdateSettings)
	protected.Get("/audit", settingsHandler.GetAuditLog)

	// Search
	protected.Get("/search", fsHandler.Search)

	// Pinned items
	protected.Get("/files/pinned", fsHandler.ListPinned)
	protected.Post("/files/pin", fsHandler.TogglePin)

	// Trash
	protected.Get("/trash", fsHandler.ListTrash)
	protected.Post("/trash/restore", fsHandler.RestoreFromTrash)
	protected.Delete("/trash", fsHandler.EmptyTrash)

	// Storage dashboard
	protected.Get("/storage/dashboard", fsHandler.StorageDashboard)

	// File operations
	protected.Post("/files/mkdir", fsHandler.CreateFolder)
	protected.Post("/files/rename", fsHandler.Rename)
	protected.Post("/files/move", fsHandler.Move)
	protected.Post("/files/copy", fsHandler.Copy)
	protected.Get("/files/download-folder", fsHandler.DownloadFolder)
	protected.Post("/files/upload", fsHandler.Upload)

	// File listing and download (with path jail)
	files := protected.Group("/files", middleware.PathJail(cfg.StorageDir))
	files.Get("/thumbnail/:checksum", fsHandler.ServeThumbnail)
	files.Get("/download/*", fsHandler.Download)
	files.Post("/zip", archiveHandler.Zip)
	files.Post("/unzip", archiveHandler.Unzip)
	files.Delete("/*", fsHandler.Delete)
	files.Get("/*", fsHandler.ListDirectory)

	// Tus resumable uploads
	tus := protected.Group("/tus")
	tus.Options("/", tusHandler.Options)
	tus.Post("/", tusHandler.Create)
	tus.Head("/:id", tusHandler.Head)
	tus.Patch("/:id", tusHandler.Patch)
	tus.Delete("/:id", tusHandler.Delete)

	// Background Tasks
	protected.Get("/tasks", archiveHandler.Tasks)

	// Sharing
	protected.Post("/share", shareHandler.CreateShare)
	protected.Get("/share", shareHandler.ListShares)
	protected.Delete("/share/:id", shareHandler.RevokeShare)

	// WebDAV Endpoint (accepts basic auth OR bearer)
	// We use All() because WebDAV uses custom HTTP methods (PROPFIND, MKCOL, LOCK, etc.)
	app.All("/webdav/*", webdavHandler.Handle)

	// Start server
	addr := fmt.Sprintf(":%s", cfg.Port)
	log.Printf("🚀 Drive server starting on %s", addr)
	log.Printf("📁 Storage: %s", cfg.StorageDir)
	log.Printf("💾 Data: %s", cfg.DataDir)

	if err := app.Listen(addr); err != nil {
		log.Fatal("Failed to start server:", err)
	}
}

func generateSecret(length int) (string, error) {
	bytes := make([]byte, length)
	if _, err := rand.Read(bytes); err != nil {
		return "", err
	}
	return hex.EncodeToString(bytes), nil
}