diff --git a/shortcuts/im/helpers_network_test.go b/shortcuts/im/helpers_network_test.go index 034393c83..a068c5b74 100644 --- a/shortcuts/im/helpers_network_test.go +++ b/shortcuts/im/helpers_network_test.go @@ -262,7 +262,7 @@ func TestDownloadIMResourceToPathSuccess(t *testing.T) { cmdutil.TestChdir(t, t.TempDir()) target := filepath.Join("nested", "resource.bin") - _, size, err := downloadIMResourceToPath(context.Background(), runtime, "om_123", "file_123", "file", target) + _, size, err := downloadIMResourceToPath(context.Background(), runtime, "om_123", "file_123", "file", target, true) if err != nil { t.Fatalf("downloadIMResourceToPath() error = %v", err) } @@ -308,7 +308,7 @@ func TestDownloadIMResourceToPathImageUsesSingleRequestWithoutRange(t *testing.T cmdutil.TestChdir(t, t.TempDir()) - gotPath, size, err := downloadIMResourceToPath(context.Background(), runtime, "om_img", "img_123", "image", "image") + gotPath, size, err := downloadIMResourceToPath(context.Background(), runtime, "om_img", "img_123", "image", "image", true) if err != nil { t.Fatalf("downloadIMResourceToPath() error = %v", err) } @@ -342,7 +342,7 @@ func TestDownloadIMResourceToPathHTTPErrorBody(t *testing.T) { cmdutil.TestChdir(t, t.TempDir()) - _, _, err := downloadIMResourceToPath(context.Background(), runtime, "om_403", "file_403", "file", "out.bin") + _, _, err := downloadIMResourceToPath(context.Background(), runtime, "om_403", "file_403", "file", "out.bin", true) if err == nil || !strings.Contains(err.Error(), "HTTP 403: denied") { t.Fatalf("downloadIMResourceToPath() error = %v", err) } @@ -372,7 +372,7 @@ func TestDownloadIMResourceToPathRetriesNetworkError(t *testing.T) { cmdutil.TestChdir(t, t.TempDir()) target := "out.bin" - _, size, err := downloadIMResourceToPath(context.Background(), runtime, "om_retry", "file_retry", "file", target) + _, size, err := downloadIMResourceToPath(context.Background(), runtime, "om_retry", "file_retry", "file", target, true) if err != nil { t.Fatalf("downloadIMResourceToPath() error = %v", err) } @@ -408,7 +408,7 @@ func TestDownloadIMResourceToPathRetrySecondAttemptSuccess(t *testing.T) { cmdutil.TestChdir(t, t.TempDir()) target := "out.bin" - _, size, err := downloadIMResourceToPath(context.Background(), runtime, "om_retry2", "file_retry2", "file", target) + _, size, err := downloadIMResourceToPath(context.Background(), runtime, "om_retry2", "file_retry2", "file", target, true) if err != nil { t.Fatalf("downloadIMResourceToPath() error = %v", err) } @@ -444,7 +444,7 @@ func TestDownloadIMResourceToPathRetryContextCanceled(t *testing.T) { cmdutil.TestChdir(t, t.TempDir()) target := "out.bin" - _, _, err := downloadIMResourceToPath(ctx, runtime, "om_cancel", "file_cancel", "file", target) + _, _, err := downloadIMResourceToPath(ctx, runtime, "om_cancel", "file_cancel", "file", target, true) if err != context.Canceled { t.Fatalf("downloadIMResourceToPath() error = %v, want context.Canceled", err) } @@ -525,7 +525,7 @@ func TestDownloadIMResourceToPathRangeDownload(t *testing.T) { cmdutil.TestChdir(t, t.TempDir()) target := filepath.Join("nested", "resource.bin") - _, size, err := downloadIMResourceToPath(context.Background(), runtime, "om_range", "file_range", "file", target) + _, size, err := downloadIMResourceToPath(context.Background(), runtime, "om_range", "file_range", "file", target, true) if err != nil { t.Fatalf("downloadIMResourceToPath() error = %v", err) } @@ -567,7 +567,7 @@ func TestDownloadIMResourceToPathInvalidContentRange(t *testing.T) { })) cmdutil.TestChdir(t, t.TempDir()) - _, _, err := downloadIMResourceToPath(context.Background(), runtime, "om_bad", "file_bad", "file", "out.bin") + _, _, err := downloadIMResourceToPath(context.Background(), runtime, "om_bad", "file_bad", "file", "out.bin", true) if err == nil || !strings.Contains(err.Error(), "invalid Content-Range header") { t.Fatalf("downloadIMResourceToPath() error = %v", err) } @@ -596,7 +596,7 @@ func TestDownloadIMResourceToPathRangeChunkFailureCleansOutput(t *testing.T) { cmdutil.TestChdir(t, t.TempDir()) target := "out.bin" - _, _, err := downloadIMResourceToPath(context.Background(), runtime, "om_miderr", "file_miderr", "file", target) + _, _, err := downloadIMResourceToPath(context.Background(), runtime, "om_miderr", "file_miderr", "file", target, true) if err == nil || !strings.Contains(err.Error(), "HTTP 500: chunk failed") { t.Fatalf("downloadIMResourceToPath() error = %v", err) } @@ -622,7 +622,7 @@ func TestDownloadIMResourceToPathRangeOverflowCleansOutput(t *testing.T) { cmdutil.TestChdir(t, t.TempDir()) target := "out.bin" - _, _, err := downloadIMResourceToPath(context.Background(), runtime, "om_overflow", "file_overflow", "file", target) + _, _, err := downloadIMResourceToPath(context.Background(), runtime, "om_overflow", "file_overflow", "file", target, true) if err == nil || !strings.Contains(err.Error(), "chunk overflow") { t.Fatalf("downloadIMResourceToPath() error = %v", err) } @@ -658,7 +658,7 @@ func TestDownloadIMResourceToPathRangeShortChunkSizeMismatch(t *testing.T) { cmdutil.TestChdir(t, t.TempDir()) - _, _, err := downloadIMResourceToPath(context.Background(), runtime, "om_short", "file_short", "file", "out.bin") + _, _, err := downloadIMResourceToPath(context.Background(), runtime, "om_short", "file_short", "file", "out.bin", true) if err == nil || !strings.Contains(err.Error(), "file size mismatch") { t.Fatalf("downloadIMResourceToPath() error = %v", err) } diff --git a/shortcuts/im/helpers_test.go b/shortcuts/im/helpers_test.go index 2596b8f0e..6de1cff1f 100644 --- a/shortcuts/im/helpers_test.go +++ b/shortcuts/im/helpers_test.go @@ -593,7 +593,7 @@ func TestDownloadIMResourceToPathHTTPClientError(t *testing.T) { return nil, errors.New("http client unavailable") })) - _, _, err := downloadIMResourceToPath(context.Background(), runtime, "om_123", "img_123", "image", "out.bin") + _, _, err := downloadIMResourceToPath(context.Background(), runtime, "om_123", "img_123", "image", "out.bin", true) if err == nil || !strings.Contains(err.Error(), "http client unavailable") { t.Fatalf("downloadIMResourceToPath() error = %v", err) } @@ -637,6 +637,68 @@ func TestParseTotalSize(t *testing.T) { } } +func TestParseContentDispositionFilename(t *testing.T) { + tests := []struct { + name string + header string + want string + }{ + {name: "empty header", header: "", want: ""}, + {name: "no filename param", header: "attachment", want: ""}, + {name: "plain filename", header: `attachment; filename="report.xlsx"`, want: "report.xlsx"}, + {name: "unquoted filename", header: `attachment; filename=report.xlsx`, want: "report.xlsx"}, + {name: "RFC 5987 UTF-8 encoded", header: `attachment; filename*=UTF-8''%E5%AD%A3%E5%BA%A6%E6%8A%A5%E5%91%8A.xlsx`, want: "季度报告.xlsx"}, + {name: "RFC 5987 takes priority over plain", header: `attachment; filename="fallback.xlsx"; filename*=UTF-8''%E5%AD%A3%E5%BA%A6%E6%8A%A5%E5%91%8A.xlsx`, want: "季度报告.xlsx"}, + {name: "path traversal stripped", header: `attachment; filename="../../etc/passwd"`, want: "passwd"}, + {name: "windows path stripped", header: `attachment; filename="C:\\Windows\\evil.exe"`, want: "evil.exe"}, + {name: "control char rejected", header: "attachment; filename=\"evil\x01file.txt\"", want: ""}, + {name: "malformed header", header: "not/valid/mime; ===", want: ""}, + {name: "whitespace trimmed", header: `attachment; filename=" report.pdf "`, want: "report.pdf"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := parseContentDispositionFilename(tt.header); got != tt.want { + t.Fatalf("parseContentDispositionFilename(%q) = %q, want %q", tt.header, got, tt.want) + } + }) + } +} + +func TestResolveIMResourceDownloadPath(t *testing.T) { + tests := []struct { + name string + safePath string + contentType string + contentDisposition string + userSpecifiedOutput bool + want string + }{ + // safePath already has extension: always return as-is + {name: "user path with ext, no CD", safePath: "out.xlsx", contentType: "application/pdf", userSpecifiedOutput: true, want: "out.xlsx"}, + {name: "user path with ext, CD present", safePath: "out.xlsx", contentDisposition: `attachment; filename="server.pdf"`, userSpecifiedOutput: true, want: "out.xlsx"}, + // No --output: use CD filename when present + {name: "default path, CD filename", safePath: "file_xxx", contentDisposition: `attachment; filename="季度报告.xlsx"`, want: "季度报告.xlsx"}, + {name: "default path, CD RFC5987", safePath: "file_xxx", contentDisposition: `attachment; filename*=UTF-8''%E5%AD%A3%E5%BA%A6%E6%8A%A5%E5%91%8A.xlsx`, want: "季度报告.xlsx"}, + {name: "default path, no CD, MIME ext", safePath: "file_xxx", contentType: "application/pdf", want: "file_xxx.pdf"}, + {name: "default path, no CD, unknown MIME", safePath: "file_xxx", contentType: "application/x-unknown", want: "file_xxx"}, + {name: "default path, CD with dir component", safePath: "downloads/file_xxx", contentDisposition: `attachment; filename="report.xlsx"`, want: "downloads/report.xlsx"}, + // User --output without extension: use CD filename's extension + {name: "user path no ext, CD with ext", safePath: "myfile", contentDisposition: `attachment; filename="server.pdf"`, userSpecifiedOutput: true, want: "myfile.pdf"}, + {name: "user path no ext, CD no ext, MIME ext", safePath: "myfile", contentDisposition: `attachment; filename="noext"`, contentType: "image/png", userSpecifiedOutput: true, want: "myfile.png"}, + {name: "user path no ext, no CD, MIME ext", safePath: "myfile", contentType: "image/jpeg", userSpecifiedOutput: true, want: "myfile.jpg"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := resolveIMResourceDownloadPath(tt.safePath, tt.contentType, tt.contentDisposition, tt.userSpecifiedOutput) + if got != tt.want { + t.Fatalf("resolveIMResourceDownloadPath() = %q, want %q", got, tt.want) + } + }) + } +} + func TestShortcuts(t *testing.T) { var commands []string for _, shortcut := range Shortcuts() { diff --git a/shortcuts/im/im_messages_resources_download.go b/shortcuts/im/im_messages_resources_download.go index 66da97636..f81e7850c 100644 --- a/shortcuts/im/im_messages_resources_download.go +++ b/shortcuts/im/im_messages_resources_download.go @@ -7,6 +7,7 @@ import ( "context" "fmt" "io" + "mime" "net/http" "path/filepath" "strconv" @@ -72,7 +73,8 @@ var ImMessagesResourcesDownload = common.Shortcut{ return output.ErrValidation("unsafe output path: %s", err) } - finalPath, sizeBytes, err := downloadIMResourceToPath(ctx, runtime, messageId, fileKey, fileType, relPath) + userSpecifiedOutput := runtime.Str("output") != "" + finalPath, sizeBytes, err := downloadIMResourceToPath(ctx, runtime, messageId, fileKey, fileType, relPath, userSpecifiedOutput) if err != nil { return err } @@ -262,18 +264,21 @@ func initialIMResourceDownloadHeaders(fileType string) map[string]string { } } -func downloadIMResourceToPath(ctx context.Context, runtime *common.RuntimeContext, messageID, fileKey, fileType, outputPath string) (string, int64, error) { +func downloadIMResourceToPath(ctx context.Context, runtime *common.RuntimeContext, messageID, fileKey, fileType, outputPath string, userSpecifiedOutput bool) (string, int64, error) { downloadResp, err := doIMResourceDownloadRequest(ctx, runtime, messageID, fileKey, fileType, initialIMResourceDownloadHeaders(fileType)) if err != nil { return "", 0, err } + if downloadResp == nil { + return "", 0, output.ErrNetwork("download failed: empty response") + } if downloadResp.StatusCode >= 400 { defer downloadResp.Body.Close() return "", 0, downloadResponseError(downloadResp) } - finalPath := resolveIMResourceDownloadPath(outputPath, downloadResp.Header.Get("Content-Type")) + finalPath := resolveIMResourceDownloadPath(outputPath, downloadResp.Header.Get("Content-Type"), downloadResp.Header.Get("Content-Disposition"), userSpecifiedOutput) var ( body io.ReadCloser @@ -316,18 +321,63 @@ func downloadIMResourceToPath(ctx context.Context, runtime *common.RuntimeContex return savedPath, result.Size(), nil } -func resolveIMResourceDownloadPath(safePath, contentType string) string { +func resolveIMResourceDownloadPath(safePath, contentType, contentDisposition string, userSpecifiedOutput bool) string { if filepath.Ext(safePath) != "" { return safePath } - mimeType := strings.Split(contentType, ";")[0] - mimeType = strings.TrimSpace(mimeType) + if cdFilename := parseContentDispositionFilename(contentDisposition); cdFilename != "" { + if !userSpecifiedOutput { + // No --output flag: use the original filename from the server. + dir := filepath.Dir(safePath) + if dir == "." { + return cdFilename + } + return filepath.Join(dir, cdFilename) + } + // User specified a path without extension: append the extension from the CD filename. + if ext := filepath.Ext(cdFilename); ext != "" { + return safePath + ext + } + } + mimeType := strings.TrimSpace(strings.Split(contentType, ";")[0]) if ext, ok := imMimeToExt[mimeType]; ok { return safePath + ext } return safePath } +// parseContentDispositionFilename extracts and sanitizes the filename from a +// Content-Disposition header. It handles RFC 5987 encoded filenames (filename*) +// with priority over plain filename via the standard mime package. +// Returns an empty string if no valid filename can be extracted. +func parseContentDispositionFilename(header string) string { + if header == "" { + return "" + } + _, params, err := mime.ParseMediaType(header) + if err != nil { + return "" + } + name := strings.TrimSpace(params["filename"]) + if name == "" { + return "" + } + // Strip any path component (Unix or Windows style) to prevent path traversal. + if i := strings.LastIndexAny(name, "/\\"); i >= 0 { + name = name[i+1:] + } + if name == "" || name == "." || name == ".." { + return "" + } + // Reject control characters (including null bytes). + for _, r := range name { + if r < 0x20 || r == 0x7f { + return "" + } + } + return name +} + func doIMResourceDownloadRequest(ctx context.Context, runtime *common.RuntimeContext, messageID, fileKey, fileType string, headers map[string]string) (*http.Response, error) { query := larkcore.QueryParams{} query.Set("type", fileType) diff --git a/skills/lark-im/references/lark-im-messages-resources-download.md b/skills/lark-im/references/lark-im-messages-resources-download.md index 5dd368d2d..cf10bdedb 100644 --- a/skills/lark-im/references/lark-im-messages-resources-download.md +++ b/skills/lark-im/references/lark-im-messages-resources-download.md @@ -34,7 +34,7 @@ lark-cli im +messages-resources-download --message-id om_xxx --file-key img_v3_x | `--message-id ` | Yes | Message ID (`om_xxx` format) | | `--file-key ` | Yes | Resource key (`img_xxx` or `file_xxx`) | | `--type ` | Yes | Resource type: `image` or `file` | -| `--output ` | No | Output path (relative paths only; `..` traversal is not allowed; defaults to `file_key` as the file name). File extension is automatically added based on Content-Type if not provided | +| `--output ` | No | Output path (relative paths only; `..` traversal is not allowed). When omitted, the server's original filename from `Content-Disposition` is used if available; otherwise defaults to `file_key`. File extension is automatically inferred from `Content-Disposition` or `Content-Type` if not provided | | `--as ` | No | Identity type: `user` (default) or `bot` | | `--dry-run` | No | Print the request only, do not execute it | @@ -51,7 +51,7 @@ When downloading large files, the command automatically uses **HTTP Range reques **Benefits:** - Reduces the impact of transient request failures during large downloads -- Automatically detects and appends correct file extension from Content-Type +- Preserves the server's original filename via `Content-Disposition` (supports RFC 5987 UTF-8 encoding); falls back to `Content-Type`-based extension inference - Validates file size integrity after download completion ## `file_key` Sources