summaryrefslogtreecommitdiff
path: root/internal/image/download.go
diff options
context:
space:
mode:
Diffstat (limited to 'internal/image/download.go')
-rw-r--r--internal/image/download.go270
1 files changed, 6 insertions, 264 deletions
diff --git a/internal/image/download.go b/internal/image/download.go
index aca32c5..42f733d 100644
--- a/internal/image/download.go
+++ b/internal/image/download.go
@@ -2,7 +2,6 @@ package image
import (
"context"
- "errors"
"fmt"
"io"
"os"
@@ -66,7 +65,8 @@ func (d *Downloader) DownloadImage(ctx context.Context, result *SearchResult, ou
_ = reader.Close()
}()
- file, parentFD, finalName, err := d.openSecureOutputFile(outputPath)
+ store := newSecureOutputStore(d.options)
+ file, parentFD, finalName, err := store.openSecureOutputFile(outputPath)
if err != nil {
return fmt.Errorf("create output file %q: %w", outputPath, err)
}
@@ -106,7 +106,7 @@ func (d *Downloader) DownloadImage(ctx context.Context, result *SearchResult, ou
if attribution := strings.TrimSpace(result.Attribution); attribution != "" {
attrRelPath := strings.TrimSuffix(finalName, filepath.Ext(finalName)) + "_attribution.txt"
- if err := d.writeSecureRelativeFile(outputPath, attrRelPath, []byte(attribution), 0o644); err != nil {
+ if err := store.writeSecureRelativeFile(outputPath, attrRelPath, []byte(attribution), 0o644); err != nil {
fmt.Fprintf(os.Stderr, "Warning: failed to save attribution: %v\n", err)
}
}
@@ -141,9 +141,10 @@ func (d *Downloader) DownloadBestMatchWithOptions(ctx context.Context, opts *Sea
return nil, "", fmt.Errorf("no images found for %q", opts.Query)
}
+ policy := newDownloadPathPolicy(d.options)
for i, result := range results {
- filename := d.generateFileName(opts.Query, &result, i)
- outputPath, err := d.resolveOutputPath(filename)
+ filename := policy.generateFileName(opts.Query, &result, i)
+ outputPath, err := policy.resolveOutputPath(filename)
if err != nil {
fmt.Fprintf(os.Stderr, "Warning: refusing unsafe output path %q: %v\n", filename, err)
continue
@@ -158,262 +159,3 @@ func (d *Downloader) DownloadBestMatchWithOptions(ctx context.Context, opts *Sea
return nil, "", fmt.Errorf("no downloadable images found for %q", opts.Query)
}
-
-func (d *Downloader) generateFileName(word string, result *SearchResult, index int) string {
- filename := ""
- if d != nil && d.options != nil {
- filename = d.options.FileNamePattern
- }
- if filename == "" {
- filename = "{word}_{source}"
- }
-
- filename = strings.ReplaceAll(filename, "{word}", sanitizeFileName(word))
- if result != nil {
- filename = strings.ReplaceAll(filename, "{source}", sanitizeFileName(result.Source))
- filename = strings.ReplaceAll(filename, "{id}", sanitizeFileName(result.ID))
- }
- filename = strings.ReplaceAll(filename, "{index}", fmt.Sprintf("%d", index))
-
- ext := ""
- if result != nil {
- ext = filepath.Ext(result.URL)
- if strings.HasPrefix(result.URL, geminiDataPrefix) {
- ext = ".png"
- } else if ext == "" || len(ext) > 5 {
- ext = ".jpg"
- }
- }
-
- if filepath.Ext(filename) == "" {
- filename += ext
- }
-
- return filename
-}
-
-func (d *Downloader) resolveOutputPath(name string) (string, error) {
- baseDir := "./images"
- if d != nil && d.options != nil && d.options.OutputDir != "" {
- baseDir = d.options.OutputDir
- }
-
- if strings.TrimSpace(baseDir) == "" {
- baseDir = "."
- }
-
- cleanBase, err := filepath.Abs(baseDir)
- if err != nil {
- return "", fmt.Errorf("resolve base dir: %w", err)
- }
-
- fullPath, err := filepath.Abs(filepath.Join(cleanBase, name))
- if err != nil {
- return "", fmt.Errorf("resolve output path: %w", err)
- }
-
- rel, err := filepath.Rel(cleanBase, fullPath)
- if err != nil {
- return "", fmt.Errorf("relativize output path: %w", err)
- }
- if rel == ".." || strings.HasPrefix(rel, ".."+string(os.PathSeparator)) {
- return "", fmt.Errorf("path escapes base dir")
- }
-
- return fullPath, nil
-}
-
-func (d *Downloader) openSecureOutputFile(outputPath string) (*os.File, int, string, error) {
- baseDir := "./images"
- if d != nil && d.options != nil && d.options.OutputDir != "" {
- baseDir = d.options.OutputDir
- }
-
- baseAbs, err := filepath.Abs(baseDir)
- if err != nil {
- return nil, 0, "", fmt.Errorf("resolve base dir: %w", err)
- }
- targetAbs, err := filepath.Abs(outputPath)
- if err != nil {
- return nil, 0, "", fmt.Errorf("resolve output path: %w", err)
- }
- relPath, err := filepath.Rel(baseAbs, targetAbs)
- if err != nil {
- return nil, 0, "", fmt.Errorf("relativize output path: %w", err)
- }
- if relPath == "." || relPath == ".." || strings.HasPrefix(relPath, ".."+string(os.PathSeparator)) {
- return nil, 0, "", fmt.Errorf("path escapes base dir")
- }
-
- relPath = filepath.Clean(relPath)
- parentRel := filepath.Dir(relPath)
- finalName := filepath.Base(relPath)
-
- parentFD, err := openDirPathNoFollow(baseAbs, parentRel, d.options != nil && d.options.CreateDir)
- if err != nil {
- return nil, 0, "", err
- }
-
- flags := unix.O_WRONLY | unix.O_CREAT | unix.O_CLOEXEC | unix.O_NOFOLLOW
- if d.options != nil && d.options.OverwriteExisting {
- flags |= unix.O_TRUNC
- } else {
- flags |= unix.O_EXCL
- }
-
- fd, err := unix.Openat(parentFD, finalName, flags, 0o644)
- if err != nil {
- _ = unix.Close(parentFD)
- return nil, 0, "", fmt.Errorf("open output file: %w", err)
- }
-
- return os.NewFile(uintptr(fd), targetAbs), parentFD, finalName, nil
-}
-
-func (d *Downloader) writeSecureRelativeFile(outputPath, relPath string, content []byte, perm os.FileMode) error {
- file, parentFD, finalName, err := d.openSecureRelativeFile(outputPath, relPath, perm)
- if err != nil {
- return err
- }
- defer func() {
- _ = unix.Close(parentFD)
- }()
- defer func() {
- _ = file.Close()
- }()
-
- if _, err := file.Write(content); err != nil {
- _ = unix.Unlinkat(parentFD, finalName, 0)
- return fmt.Errorf("write attribution file: %w", err)
- }
- if err := file.Sync(); err != nil {
- _ = unix.Unlinkat(parentFD, finalName, 0)
- return fmt.Errorf("sync attribution file: %w", err)
- }
- return nil
-}
-
-func (d *Downloader) openSecureRelativeFile(outputPath, relPath string, perm os.FileMode) (*os.File, int, string, error) {
- baseDir := "./images"
- if d != nil && d.options != nil && d.options.OutputDir != "" {
- baseDir = d.options.OutputDir
- }
-
- baseAbs, err := filepath.Abs(baseDir)
- if err != nil {
- return nil, 0, "", fmt.Errorf("resolve base dir: %w", err)
- }
- targetAbs, err := filepath.Abs(outputPath)
- if err != nil {
- return nil, 0, "", fmt.Errorf("resolve output path: %w", err)
- }
- baseRelPath, err := filepath.Rel(baseAbs, targetAbs)
- if err != nil {
- return nil, 0, "", fmt.Errorf("relativize output path: %w", err)
- }
- if baseRelPath == "." || baseRelPath == ".." || strings.HasPrefix(baseRelPath, ".."+string(os.PathSeparator)) {
- return nil, 0, "", fmt.Errorf("path escapes base dir")
- }
-
- parentRel := filepath.Dir(filepath.Clean(baseRelPath))
- parentFD, err := openDirPathNoFollow(baseAbs, parentRel, true)
- if err != nil {
- return nil, 0, "", err
- }
-
- flags := unix.O_WRONLY | unix.O_CREAT | unix.O_CLOEXEC | unix.O_NOFOLLOW | unix.O_TRUNC
- fd, err := unix.Openat(parentFD, relPath, flags, uint32(perm.Perm()))
- if err != nil {
- _ = unix.Close(parentFD)
- return nil, 0, "", fmt.Errorf("open attribution file: %w", err)
- }
-
- return os.NewFile(uintptr(fd), targetAbs), parentFD, relPath, nil
-}
-
-func openDirPathNoFollow(baseAbs, relPath string, createDirs bool) (int, error) {
- dirFD, err := openPathDirNoFollow(baseAbs, createDirs)
- if err != nil {
- return 0, err
- }
-
- relPath = filepath.Clean(relPath)
- if relPath == "." {
- return dirFD, nil
- }
-
- parts := strings.Split(relPath, string(os.PathSeparator))
- for _, part := range parts {
- if part == "" || part == "." {
- continue
- }
- if createDirs {
- if err := unix.Mkdirat(dirFD, part, 0o755); err != nil && !errors.Is(err, unix.EEXIST) {
- _ = unix.Close(dirFD)
- return 0, fmt.Errorf("create dir %q: %w", part, err)
- }
- }
-
- nextFD, err := unix.Openat(dirFD, part, unix.O_RDONLY|unix.O_DIRECTORY|unix.O_NOFOLLOW|unix.O_CLOEXEC, 0)
- if err != nil {
- _ = unix.Close(dirFD)
- return 0, fmt.Errorf("open dir %q: %w", part, err)
- }
- _ = unix.Close(dirFD)
- dirFD = nextFD
- }
-
- return dirFD, nil
-}
-
-func openPathDirNoFollow(absPath string, createDirs bool) (int, error) {
- absPath = filepath.Clean(absPath)
- if absPath == string(os.PathSeparator) {
- return unix.Open(absPath, unix.O_RDONLY|unix.O_DIRECTORY|unix.O_NOFOLLOW|unix.O_CLOEXEC, 0)
- }
-
- parentAbs := filepath.Dir(absPath)
- leaf := filepath.Base(absPath)
- parentFD, err := openPathDirNoFollow(parentAbs, createDirs)
- if err != nil {
- return 0, err
- }
-
- if createDirs {
- if err := unix.Mkdirat(parentFD, leaf, 0o755); err != nil && !errors.Is(err, unix.EEXIST) {
- _ = unix.Close(parentFD)
- return 0, fmt.Errorf("create dir %q: %w", leaf, err)
- }
- }
-
- dirFD, err := unix.Openat(parentFD, leaf, unix.O_RDONLY|unix.O_DIRECTORY|unix.O_NOFOLLOW|unix.O_CLOEXEC, 0)
- _ = unix.Close(parentFD)
- if err != nil {
- return 0, fmt.Errorf("open dir %q: %w", leaf, err)
- }
-
- return dirFD, nil
-}
-
-func sanitizeFileName(name string) string {
- replacer := strings.NewReplacer(
- "/", "_",
- "\\", "_",
- ":", "_",
- "*", "_",
- "?", "_",
- "\"", "_",
- "<", "_",
- ">", "_",
- "|", "_",
- " ", "_",
- ".", "_",
- )
-
- sanitized := replacer.Replace(name)
- if len(sanitized) > 50 {
- sanitized = sanitized[:50]
- }
-
- return sanitized
-}