feat: add support for file downloads from binary response endpoints

Author: stainless-app[bot]

pkg/cmd/cmdutil.go

@@ -6,11 +6,13 @@ import (
 	"fmt"
 	"io"
 	"log"
+	"mime"
 	"net/http"
 	"net/http/httputil"
 	"os"
 	"os/exec"
 	"os/signal"
+	"path/filepath"
 	"strings"
 	"syscall"
 
@@ -154,6 +156,108 @@ func streamToStdout(generateOutput func(w *os.File) error) error {
 	return err
 }
 
+func writeBinaryResponse(response *http.Response, outfile string) (string, error) {
+	defer response.Body.Close()
+	body, err := io.ReadAll(response.Body)
+	if err != nil {
+		return "", err
+	}
+	switch outfile {
+	case "-", "/dev/stdout":
+		_, err := os.Stdout.Write(body)
+		return "", err
+	case "":
+		// If output file is unspecified, then print to stdout for plain text or
+		// if stdout is not a terminal:
+		if !isTerminal(os.Stdout) || isUTF8TextFile(body) {
+			_, err := os.Stdout.Write(body)
+			return "", err
+		}
+
+		// If response has a suggested filename in the content-disposition
+		// header, then use that (with an optional suffix to ensure uniqueness):
+		file, err := createDownloadFile(response, body)
+		if err != nil {
+			return "", err
+		}
+		defer file.Close()
+		if _, err := file.Write(body); err != nil {
+			return "", err
+		}
+		return fmt.Sprintf("Wrote output to: %s", file.Name()), nil
+	default:
+		if err := os.WriteFile(outfile, body, 0644); err != nil {
+			return "", err
+		}
+		return fmt.Sprintf("Wrote output to: %s", outfile), nil
+	}
+}
+
+// Return a writable file handle to a new file, which attempts to choose a good filename
+// based on the Content-Disposition header or sniffing the MIME filetype of the response.
+func createDownloadFile(response *http.Response, data []byte) (*os.File, error) {
+	filename := "file"
+	// If the header provided an output filename, use that
+	disp := response.Header.Get("Content-Disposition")
+	_, params, err := mime.ParseMediaType(disp)
+	if err == nil {
+		if dispFilename, ok := params["filename"]; ok {
+			// Only use the last path component to prevent directory traversal
+			filename = filepath.Base(dispFilename)
+			// Try to create the file with exclusive flag to avoid race conditions
+			file, err := os.OpenFile(filename, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0644)
+			if err == nil {
+				return file, nil
+			}
+		}
+	}
+
+	// If file already exists, create a unique filename using CreateTemp
+	ext := filepath.Ext(filename)
+	if ext == "" {
+		ext = guessExtension(data)
+	}
+	base := strings.TrimSuffix(filename, ext)
+	return os.CreateTemp(".", base+"-*"+ext)
+}
+
+func guessExtension(data []byte) string {
+	ct := http.DetectContentType(data)
+
+	// Prefer common extensions over obscure ones
+	switch ct {
+	case "application/gzip":
+		return ".gz"
+	case "application/pdf":
+		return ".pdf"
+	case "application/zip":
+		return ".zip"
+	case "audio/mpeg":
+		return ".mp3"
+	case "image/bmp":
+		return ".bmp"
+	case "image/gif":
+		return ".gif"
+	case "image/jpeg":
+		return ".jpg"
+	case "image/png":
+		return ".png"
+	case "image/webp":
+		return ".webp"
+	case "video/mp4":
+		return ".mp4"
+	}
+
+	exts, err := mime.ExtensionsByType(ct)
+	if err == nil && len(exts) > 0 {
+		return exts[0]
+	} else if isUTF8TextFile(data) {
+		return ".txt"
+	} else {
+		return ".bin"
+	}
+}
+
 func shouldUseColors(w io.Writer) bool {
 	force, ok := os.LookupEnv("FORCE_COLOR")
 	if ok {

pkg/cmd/cmdutil_test.go

@@ -1,8 +1,15 @@
 package cmd
 
 import (
+	"bytes"
+	"io"
+	"net/http"
 	"os"
+	"path/filepath"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestStreamOutput(t *testing.T) {
@@ -15,3 +22,106 @@ func TestStreamOutput(t *testing.T) {
 		t.Errorf("streamOutput failed: %v", err)
 	}
 }
+
+func TestWriteBinaryResponse(t *testing.T) {
+	t.Run("write to explicit file", func(t *testing.T) {
+		tmpDir := t.TempDir()
+		outfile := tmpDir + "/output.txt"
+		body := []byte("test content")
+		resp := &http.Response{
+			Body: io.NopCloser(bytes.NewReader(body)),
+		}
+
+		msg, err := writeBinaryResponse(resp, outfile)
+
+		require.NoError(t, err)
+		assert.Contains(t, msg, outfile)
+
+		content, err := os.ReadFile(outfile)
+		require.NoError(t, err)
+		assert.Equal(t, body, content)
+	})
+
+	t.Run("write to stdout", func(t *testing.T) {
+		oldStdout := os.Stdout
+		r, w, _ := os.Pipe()
+		os.Stdout = w
+
+		body := []byte("stdout content")
+		resp := &http.Response{
+			Body: io.NopCloser(bytes.NewReader(body)),
+		}
+		msg, err := writeBinaryResponse(resp, "-")
+
+		w.Close()
+		os.Stdout = oldStdout
+
+		require.NoError(t, err)
+		assert.Empty(t, msg)
+
+		var buf bytes.Buffer
+		_, _ = buf.ReadFrom(r)
+		assert.Equal(t, body, buf.Bytes())
+	})
+}
+
+func TestCreateDownloadFile(t *testing.T) {
+	t.Run("creates file with filename from header", func(t *testing.T) {
+		tmpDir := t.TempDir()
+		oldWd, _ := os.Getwd()
+		os.Chdir(tmpDir)
+		defer os.Chdir(oldWd)
+
+		resp := &http.Response{
+			Header: http.Header{
+				"Content-Disposition": []string{`attachment; filename="test.txt"`},
+			},
+		}
+		file, err := createDownloadFile(resp, []byte("test content"))
+		require.NoError(t, err)
+		defer file.Close()
+		assert.Equal(t, "test.txt", filepath.Base(file.Name()))
+
+		// Create a second file with the same name to ensure it doesn't clobber the first
+		resp2 := &http.Response{
+			Header: http.Header{
+				"Content-Disposition": []string{`attachment; filename="test.txt"`},
+			},
+		}
+		file2, err := createDownloadFile(resp2, []byte("second content"))
+		require.NoError(t, err)
+		defer file2.Close()
+		assert.NotEqual(t, file.Name(), file2.Name(), "second file should have a different name")
+		assert.Contains(t, filepath.Base(file2.Name()), "test")
+	})
+
+	t.Run("creates temp file when no header", func(t *testing.T) {
+		tmpDir := t.TempDir()
+		oldWd, _ := os.Getwd()
+		os.Chdir(tmpDir)
+		defer os.Chdir(oldWd)
+
+		resp := &http.Response{Header: http.Header{}}
+		file, err := createDownloadFile(resp, []byte("test content"))
+		require.NoError(t, err)
+		defer file.Close()
+		assert.Contains(t, filepath.Base(file.Name()), "file-")
+	})
+
+	t.Run("prevents directory traversal", func(t *testing.T) {
+		tmpDir := t.TempDir()
+		oldWd, _ := os.Getwd()
+		os.Chdir(tmpDir)
+		defer os.Chdir(oldWd)
+
+		resp := &http.Response{
+			Header: http.Header{
+				"Content-Disposition": []string{`attachment; filename="../../../etc/passwd"`},
+			},
+		}
+		file, err := createDownloadFile(resp, []byte("test content"))
+		require.NoError(t, err)
+		defer file.Close()
+		assert.Equal(t, "passwd", filepath.Base(file.Name()))
+	})
+}