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(), "_") }