package main import ( "database/sql" "encoding/json" "fmt" "log" "net/http" "os" "path/filepath" "strings" "github.com/joho/godotenv" _ "github.com/mattn/go-sqlite3" "github.com/google/uuid" ) type Config struct { Port string URLPrefix string } type Image struct { ID int Name string Delivered bool Score int SecretKey string User string } var db *sql.DB var config Config func main() { // Load environment variables err := godotenv.Load() if err != nil { log.Fatal("Error loading .env file") } config = Config{ Port: os.Getenv("PORT"), URLPrefix: os.Getenv("URL_PREFIX"), } if len(os.Args) < 2 { log.Fatal("Usage: ./image_manager [save_name|serve]") } command := os.Args[1] // Initialize database initDB() switch command { case "save_name": saveImageNames() case "serve": serveAPI() default: log.Fatal("Unknown command. Use 'save_name' or 'serve'") } } func initDB() { var err error db, err = sql.Open("sqlite3", "./image_asset.sqlite") if err != nil { log.Fatal(err) } // Create table if not exists _, err = db.Exec(` CREATE TABLE IF NOT EXISTS images ( id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT NOT NULL UNIQUE, delivered INTEGER DEFAULT 0, score INTEGER DEFAULT 0, secret_key TEXT, user TEXT ) `) if err != nil { log.Fatal(err) } } func saveImageNames() { // Get all files in image_asset directory files, err := os.ReadDir("./image_asset") if err != nil { log.Fatal(err) } // Insert image names into database for _, file := range files { if !file.IsDir() { // Check if the file is an image (simple extension check) ext := strings.ToLower(filepath.Ext(file.Name())) if ext == ".jpg" || ext == ".jpeg" || ext == ".png" || ext == ".gif" { _, err := db.Exec("INSERT OR IGNORE INTO images (name) VALUES (?)", file.Name()) if err != nil { log.Printf("Error inserting %s: %v", file.Name(), err) } else { log.Printf("Saved: %s", file.Name()) } } } } } func serveAPI() { // Create a new mux router router := http.NewServeMux() router.HandleFunc("/get-image", getImageHandler) router.HandleFunc("/save-score", saveScoreHandler) log.Printf("Server starting on port %s", config.Port) log.Fatal(http.ListenAndServe(":"+config.Port, router)) } func getImageHandler(w http.ResponseWriter, r *http.Request) { if r.Method != http.MethodGet { http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) return } // Get a random undelivered image var img Image err := db.QueryRow(` SELECT id, name FROM images WHERE delivered = 0 ORDER BY RANDOM() LIMIT 1 `).Scan(&img.ID, &img.Name) if err != nil { if err == sql.ErrNoRows { http.Error(w, "No images available", http.StatusNotFound) return } http.Error(w, err.Error(), http.StatusInternalServerError) return } // Generate secret key secretKey := uuid.New().String() // Update the image with delivered status and secret key _, err = db.Exec(` UPDATE images SET delivered = 1, secret_key = ? WHERE id = ? `, secretKey, img.ID) if err != nil { http.Error(w, err.Error(), http.StatusInternalServerError) return } // Prepare response response := map[string]interface{}{ "url": fmt.Sprintf("%s/%s", config.URLPrefix, img.Name), "secretKey": secretKey, } w.Header().Set("Content-Type", "application/json") json.NewEncoder(w).Encode(response) } func saveScoreHandler(w http.ResponseWriter, r *http.Request) { if r.Method != http.MethodPost { http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) return } type Request struct { FileName string `json:"fileName"` Score int `json:"score"` User string `json:"user,omitempty"` SecretKey string `json:"secretKey"` } var req Request err := json.NewDecoder(r.Body).Decode(&req) if err != nil { http.Error(w, err.Error(), http.StatusBadRequest) return } // Validate request if req.FileName == "" || req.SecretKey == "" { http.Error(w, "fileName and secretKey are required", http.StatusBadRequest) return } // Verify secret key and update score var img Image err = db.QueryRow(` SELECT id FROM images WHERE name = ? AND secret_key = ? `, req.FileName, req.SecretKey).Scan(&img.ID) if err != nil { if err == sql.ErrNoRows { http.Error(w, "Invalid fileName or secretKey", http.StatusUnauthorized) return } http.Error(w, err.Error(), http.StatusInternalServerError) return } // Update score and user _, err = db.Exec(` UPDATE images SET score = ?, user = ? WHERE id = ? `, req.Score, req.User, img.ID) if err != nil { http.Error(w, err.Error(), http.StatusInternalServerError) return } w.WriteHeader(http.StatusOK) w.Write([]byte("Score saved successfully")) }