gaslight/internal/capture/images.go

package capture

import (
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"mime"
	"net/http"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"sort"
	"strings"
	"sync/atomic"
	"time"

	"git.gulenok.ru/greenhaze/gaslight/internal/domain"
)

const byURLDir = "by-url"

type ImageSaver struct {
	outputDir string
	counter   uint64
}

func NewImageSaver(outputDir string) *ImageSaver {
	return &ImageSaver{outputDir: outputDir}
}

func (s *ImageSaver) Save(resp *http.Response, body []byte) error {
	if resp == nil || resp.Request == nil {
		return fmt.Errorf("nil response or request")
	}

	n := atomic.AddUint64(&s.counter, 1)
	ts := time.Now().UTC().Format("20060102T150405.000000000Z")
	host := sanitizeName(domain.HostWithoutPort(resp.Request.Host))
	ext := imageExtension(resp.Header.Get("Content-Type"), resp.Request.URL.Path)
	name := fmt.Sprintf("%s_%06d_%s%s", ts, n, host, ext)
	if err := os.WriteFile(filepath.Join(s.outputDir, name), body, 0o644); err != nil {
		return err
	}

	relPath, err := relativeImagePath(resp.Request.Host, resp.Request.URL.Path, resp.Request.URL.RawQuery, resp.Header.Get("Content-Type"))
	if err != nil {
		return err
	}
	localPath := filepath.Join(s.outputDir, byURLDir, relPath)
	if err := os.MkdirAll(filepath.Dir(localPath), 0o755); err != nil {
		return err
	}
	return os.WriteFile(localPath, body, 0o644)
}

func (s *ImageSaver) LocalPathByRelativeURL(host, relativeURL string) (string, error) {
	if strings.TrimSpace(relativeURL) == "" {
		return "", fmt.Errorf("relative url is empty")
	}

	parsed, err := url.Parse(relativeURL)
	if err != nil {
		return "", err
	}

	resolvedHost := host
	if parsed.Host != "" {
		resolvedHost = parsed.Host
	}

	resolvedPath := parsed.Path
	if resolvedPath == "" {
		resolvedPath = "/"
	}

	localPath, err := resolveByRelativeURL(s.outputDir, resolvedHost, resolvedPath, parsed.RawQuery)
	if err != nil {
		return "", err
	}
	return localPath, nil
}

func (s *ImageSaver) ReadByRelativeURL(host, relativeURL string) ([]byte, string, error) {
	localPath, err := s.LocalPathByRelativeURL(host, relativeURL)
	if err != nil {
		return nil, "", err
	}
	data, err := os.ReadFile(localPath)
	if err != nil {
		return nil, "", err
	}
	return data, localPath, nil
}

func IsImageResponse(contentType, path string) bool {
	if strings.HasPrefix(strings.ToLower(contentType), "image/") {
		return true
	}
	switch strings.ToLower(filepath.Ext(path)) {
	case ".jpg", ".jpeg", ".png", ".gif", ".bmp", ".webp", ".svg", ".ico", ".avif":
		return true
	default:
		return false
	}
}

func relativeImagePath(host, urlPath, rawQuery, contentType string) (string, error) {
	hostDir := sanitizeName(domain.HostWithoutPort(host))
	dirSegments, baseName, baseExt := splitAndSanitizeURLPath(urlPath)
	ext := baseExt
	if ext == "" {
		ext = imageExtension(contentType, urlPath)
	}
	suffix := querySuffix(rawQuery)
	fileName := baseName + suffix + ext
	parts := append([]string{hostDir}, dirSegments...)
	parts = append(parts, fileName)
	return filepath.Join(parts...), nil
}

func resolveByRelativeURL(outputDir, host, urlPath, rawQuery string) (string, error) {
	hostDir := sanitizeName(domain.HostWithoutPort(host))
	dirSegments, baseName, baseExt := splitAndSanitizeURLPath(urlPath)
	suffix := querySuffix(rawQuery)
	baseDir := filepath.Join(append([]string{outputDir, byURLDir, hostDir}, dirSegments...)...)
	prefix := baseName + suffix

	if baseExt != "" {
		candidate := filepath.Join(baseDir, prefix+strings.ToLower(baseExt))
		if fileExists(candidate) {
			return candidate, nil
		}
	} else {
		candidate := filepath.Join(baseDir, prefix+".img")
		if fileExists(candidate) {
			return candidate, nil
		}
	}

	matches, err := findMatchingFiles(baseDir, prefix)
	if err != nil {
		return "", err
	}
	if len(matches) == 0 {
		return "", os.ErrNotExist
	}
	return matches[0], nil
}

func findMatchingFiles(dir, prefix string) ([]string, error) {
	entries, err := os.ReadDir(dir)
	if err != nil {
		return nil, err
	}
	var matches []string
	for _, entry := range entries {
		if entry.IsDir() {
			continue
		}
		name := entry.Name()
		if strings.HasPrefix(name, prefix+".") {
			matches = append(matches, filepath.Join(dir, name))
		}
	}
	sort.Strings(matches)
	return matches, nil
}

func fileExists(path string) bool {
	info, err := os.Stat(path)
	return err == nil && !info.IsDir()
}

func splitAndSanitizeURLPath(urlPath string) ([]string, string, string) {
	cleanPath := path.Clean("/" + strings.TrimSpace(urlPath))
	if cleanPath == "/" {
		cleanPath = "/index"
	}
	dir := path.Dir(cleanPath)
	base := path.Base(cleanPath)
	if base == "." || base == "/" || base == "" {
		base = "index"
	}
	ext := strings.ToLower(path.Ext(base))
	name := strings.TrimSuffix(base, ext)
	name = sanitizePathSegment(name)
	if name == "" {
		name = "image"
	}

	var segments []string
	if dir != "." && dir != "/" {
		for _, seg := range strings.Split(strings.TrimPrefix(dir, "/"), "/") {
			if seg == "" || seg == "." || seg == ".." {
				continue
			}
			cleanSeg := sanitizePathSegment(seg)
			if cleanSeg == "" {
				cleanSeg = "segment"
			}
			segments = append(segments, cleanSeg)
		}
	}
	return segments, name, ext
}

func querySuffix(rawQuery string) string {
	rawQuery = strings.TrimSpace(rawQuery)
	if rawQuery == "" {
		return ""
	}
	sum := sha256.Sum256([]byte(rawQuery))
	return "__q_" + hex.EncodeToString(sum[:8])
}

func imageExtension(contentType, requestPath string) string {
	if ext := strings.ToLower(path.Ext(requestPath)); ext != "" && ext != "." {
		return ext
	}

	base := strings.ToLower(strings.TrimSpace(strings.Split(contentType, ";")[0]))
	if base != "" {
		if exts, err := mime.ExtensionsByType(base); err == nil && len(exts) > 0 {
			return exts[0]
		}
	}
	return ".img"
}

func sanitizeName(s string) string {
	s = strings.ToLower(strings.TrimSpace(s))
	if s == "" {
		return "unknown-host"
	}

	var b strings.Builder
	for _, r := range s {
		switch {
		case r >= 'a' && r <= 'z':
			b.WriteRune(r)
		case r >= '0' && r <= '9':
			b.WriteRune(r)
		case r == '.', r == '-', r == '_':
			b.WriteRune(r)
		default:
			b.WriteByte('_')
		}
	}
	return b.String()
}

func sanitizePathSegment(s string) string {
	s = strings.ToLower(strings.TrimSpace(s))
	if s == "" {
		return ""
	}

	var b strings.Builder
	for _, r := range s {
		switch {
		case r >= 'a' && r <= 'z':
			b.WriteRune(r)
		case r >= '0' && r <= '9':
			b.WriteRune(r)
		case r == '.', r == '-', r == '_':
			b.WriteRune(r)
		default:
			b.WriteByte('_')
		}
	}
	return strings.Trim(b.String(), "_")
}