singlefile-receiver/main.go

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"strings"
)

type Config struct {
	Port        string `json:"port"`
	AuthToken   string `json:"authToken"`
	SaveBaseDir string `json:"saveBaseDir"`
	DataField   string `json:"dataField"`
	URLField    string `json:"urlField"`
}

func (c *Config) Default() Config {
	return Config{
		Port:        "3600",
		AuthToken:   "change-me",
		SaveBaseDir: "./archives",
		DataField:   "data",
		URLField:    "url",
	}
}

func main() {
	configFile := flag.String("config", "", "Path to config file (JSON)")
	flag.Parse()

	config := loadConfig(*configFile)

	if err := os.MkdirAll(config.SaveBaseDir, 0755); err != nil {
		log.Fatalf("Failed to create archives directory: %v", err)
	}

	http.HandleFunc("/archive", func(w http.ResponseWriter, r *http.Request) {
		handleArchive(w, r, config)
	})

	log.Printf("Server starting on port %s", config.Port)
	log.Fatal(http.ListenAndServe(":"+config.Port, nil))
}

func loadConfig(configFile string) Config {
	var config Config
	config = config.Default()

	if configFile != "" {
		file, err := os.Open(configFile)
		if err != nil {
			log.Fatalf("Failed to open config file: %v", err)
		}
		defer func(file *os.File) {
			err := file.Close()
			if err != nil {
				log.Printf("Failed to close config file: %v", err)
			}
		}(file)

		decoder := json.NewDecoder(file)
		if err := decoder.Decode(&config); err != nil {
			log.Fatalf("Failed to parse config file: %v", err)
		}
	}

	log.Printf("Using configuration: Port=%s, SaveBaseDir=%s", config.Port, config.SaveBaseDir)
	return config
}

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

	authHeader := r.Header.Get("Authorization")
	if !validateBearerToken(authHeader, config.AuthToken) {
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}

	err := r.ParseMultipartForm(10 << 20) // 10 MB max memory
	if err != nil {
		http.Error(w, "Failed to parse multipart form: "+err.Error(), http.StatusBadRequest)
		return
	}

	url := r.FormValue(config.URLField)
	if url == "" {
		http.Error(w, "URL field is required", http.StatusBadRequest)
		return
	}

	file, fileHeader, err := r.FormFile(config.DataField)
	if err != nil {
		http.Error(w, "Failed to get file from form: "+err.Error(), http.StatusBadRequest)
		return
	}
	defer func(file io.ReadCloser) {
		err := file.Close()
		if err != nil {
			log.Println("Failed to close multipart form file")
		}
	}(file)

	contentBytes, err := io.ReadAll(file)
	if err != nil {
		http.Error(w, "Failed to read file content: "+err.Error(), http.StatusInternalServerError)
		return
	}

	if len(contentBytes) == 0 {
		http.Error(w, "Uploaded file is empty", http.StatusBadRequest)
		return
	}

	var filePath string

	if fileHeader != nil && fileHeader.Filename != "" {
		filename := sanitizeFilename(fileHeader.Filename)
		filePath = filepath.Join(config.SaveBaseDir, filename)

		if !strings.HasSuffix(strings.ToLower(filePath), ".html") {
			filePath += ".html"
		}
	} else {
		filePath, err = getSanitizedFilePath(url, config.SaveBaseDir)
		if err != nil {
			http.Error(w, "Failed to process URL: "+err.Error(), http.StatusInternalServerError)
			return
		}
	}

	if err := os.WriteFile(filePath, contentBytes, 0644); err != nil {
		http.Error(w, "Failed to save file: "+err.Error(), http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusCreated)

	response := map[string]string{
		"message": "Page archived successfully",
		"path":    filePath,
	}

	if err := json.NewEncoder(w).Encode(response); err != nil {
		log.Printf("Error encoding response: %v", err)
	}
}

func validateBearerToken(authHeader string, authToken string) bool {
	if !strings.HasPrefix(authHeader, "Bearer ") {
		return false
	}

	token := strings.TrimPrefix(authHeader, "Bearer ")
	return token == authToken
}

func getSanitizedFilePath(url string, saveBaseDir string) (string, error) {
	url = regexp.MustCompile(`^https?://`).ReplaceAllString(url, "")
	url = regexp.MustCompile(`[/<>:"\\|?*]`).ReplaceAllString(url, "_")
	fileName := url

	if err := os.MkdirAll(saveBaseDir, 0755); err != nil {
		return "", fmt.Errorf("failed to create directory: %w", err)
	}

	fullPath := filepath.Join(saveBaseDir, fileName)
	if !strings.HasSuffix(fullPath, ".html") {
		fullPath += ".html"
	}

	return fullPath, nil
}

func sanitizeFilename(filename string) string {
	return regexp.MustCompile(`[/<>:"\\|?*]`).ReplaceAllString(filename, "_")
}