odroid-image-server/main.go

package main

import (
	"database/sql"
	"flag"
	"fmt"
	"html/template"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/fsnotify/fsnotify"
	_ "github.com/mattn/go-sqlite3" // SQLite driver
)

// --- Settings (defined as command-line flags) ---
var (
	imageRootPath *string
	templatesPath *string
	port          *string
	dbFile        *string
)

// FileInfo struct to hold data for the template
type FileInfo struct {
	Board       string
	Release     string
	Variant     string
	DownloadURL string
	FileName    string
	Checksum    string
	Count       int64
}

// AppState struct for thread-safe data storage
type AppState struct {
	sync.RWMutex
	imagesByBoard map[string][]FileInfo
}

var (
	appState = AppState{imagesByBoard: make(map[string][]FileInfo)}
	db       *sql.DB // Database connection pool
)

// initDB initializes the database and creates the table if it doesn't exist
func initDB() {
	var err error
	db, err = sql.Open("sqlite3", *dbFile)
	if err != nil {
		log.Fatalf("Failed to connect to DB: %v", err)
	}

	createTableSQL := `
	CREATE TABLE IF NOT EXISTS downloads (
		path TEXT PRIMARY KEY,
		count INTEGER NOT NULL DEFAULT 0
	);`
	if _, err := db.Exec(createTableSQL); err != nil {
		log.Fatalf("Failed to create table: %v", err)
	}
	log.Printf("Successfully connected to database '%s'.", *dbFile)
}

// refreshImageList scans the image directory and updates the file list
func refreshImageList() {
	log.Println("Refreshing image list...")
	newImages := make(map[string][]FileInfo)

	err := filepath.Walk(*imageRootPath, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if !info.IsDir() && strings.HasSuffix(info.Name(), ".img.xz") {
			relPath, _ := filepath.Rel(*imageRootPath, path)
			parts := strings.Split(filepath.ToSlash(relPath), "/")

			if len(parts) >= 4 {
				file := FileInfo{
					Board:       strings.ToUpper(strings.Replace(parts[0], "odroid", "ODROID-", 1)),
					Release:     fmt.Sprintf("Ubuntu %s (%s)", parseVersion(parts[1]), strings.Title(parts[1])),
					Variant:     strings.Title(strings.Replace(parts[2], "-", " ", -1)),
					DownloadURL: "/download/" + filepath.ToSlash(relPath),
					FileName:    info.Name(),
				}

				// Read MD5 checksum file
				checksumPath := strings.Replace(path, ".img.xz", ".img.xz.md5sum", 1)
				if content, err := os.ReadFile(checksumPath); err == nil {
					file.Checksum = strings.Fields(string(content))[0]
				} else {
					file.Checksum = "N/A"
				}

				// Query download count from DB
				err := db.QueryRow("SELECT count FROM downloads WHERE path = ?", file.DownloadURL).Scan(&file.Count)
				if err != nil && err != sql.ErrNoRows {
					log.Printf("Failed to query count for %s: %v", file.DownloadURL, err)
				}

				newImages[file.Board] = append(newImages[file.Board], file)
			}
		}
		return nil
	})

	if err != nil {
		log.Printf("Error while scanning directory: %v", err)
		return
	}

	for board := range newImages {
		sort.Slice(newImages[board], func(i, j int) bool {
			if newImages[board][i].Release != newImages[board][j].Release {
				return newImages[board][i].Release < newImages[board][j].Release
			}
			return newImages[board][i].Variant < newImages[board][j].Variant
		})
	}

	appState.Lock()
	appState.imagesByBoard = newImages
	appState.Unlock()
	log.Println("Image list refresh complete.")
}

// watchFiles detects file system changes and triggers a list refresh
func watchFiles() {
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		log.Fatalf("Failed to create file watcher: %v", err)
	}
	defer watcher.Close()

	if err := filepath.Walk(*imageRootPath, func(path string, info os.FileInfo, err error) error {
		if info.IsDir() {
			return watcher.Add(path)
		}
		return nil
	}); err != nil {
		log.Fatalf("Failed to add directories to watcher: %v", err)
	}

	var timer *time.Timer
	for {
		select {
		case event, ok := <-watcher.Events:
			if !ok {
				return
			}
			if event.Op&(fsnotify.Create|fsnotify.Remove|fsnotify.Rename) != 0 {
				log.Printf("Detected file system change: %s", event.Name)
				if timer != nil {
					timer.Stop()
				}
				timer = time.AfterFunc(2*time.Second, func() {
					refreshImageList()
					if err := watcher.Add(event.Name); err != nil && !os.IsNotExist(err) {
						log.Printf("Failed to add new directory to watcher: %v", err)
					}
				})
			}
		case err, ok := <-watcher.Errors:
			if !ok {
				return
			}
			log.Println("File watcher error:", err)
		}
	}
}

func parseVersion(release string) string {
	switch release {
	case "focal":
		return "20.04"
	case "jammy":
		return "22.04"
	case "noble":
		return "24.04"
	default:
		return release
	}
}

func main() {
	// 1. Define command-line flags
	port = flag.String("port", "8080", "Web server port number")
	imageRootPath = flag.String("imageDir", "public/image/monthly", "Image file root directory")
	templatesPath = flag.String("templateDir", "templates", "HTML template directory")
	dbFile = flag.String("dbFile", "metadata.db", "SQLite database file path")
	flag.Parse()

	// 2. Initialization
	initDB()
	defer db.Close()
	refreshImageList()

	// 3. Start the file watcher goroutine
	go watchFiles()

	// 4. Set up web server handlers
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		appState.RLock()
		boards := make([]string, 0, len(appState.imagesByBoard))
		for k := range appState.imagesByBoard {
			boards = append(boards, k)
		}
		sort.Strings(boards)

		data := struct {
			Title, PageHeader string
			Boards            []string
			Images            map[string][]FileInfo
		}{
			Title:      "ODROID OS Image Downloads",
			PageHeader: "ODROID OS Image Downloads",
			Boards:     boards,
			Images:     appState.imagesByBoard,
		}
		appState.RUnlock()

		tmpl, err := template.ParseFiles(filepath.Join(*templatesPath, "index.html"))
		if err != nil {
			http.Error(w, "Could not load template.", http.StatusInternalServerError)
			log.Printf("Template parsing error: %v", err)
			return
		}
		tmpl.Execute(w, data)
	})

	http.HandleFunc("/download/", func(w http.ResponseWriter, r *http.Request) {
		// Increment download count (UPSERT)
		upsertSQL := `
		INSERT INTO downloads (path, count) VALUES (?, 1)
		ON CONFLICT(path) DO UPDATE SET count = count + 1;`

		if _, err := db.Exec(upsertSQL, r.URL.Path); err != nil {
			log.Printf("Failed to update DB count for %s: %v", r.URL.Path, err)
		}

		filePath := filepath.Join(*imageRootPath, strings.TrimPrefix(r.URL.Path, "/download/"))
		http.ServeFile(w, r, filePath)
	})

	// 5. Start the web server
	log.Printf("Starting server at http://localhost:%s", *port)
	if err := http.ListenAndServe(":"+*port, nil); err != nil {
		log.Fatalf("Failed to start server: %v", err)
	}
}