diff --git a/pkg/cyberhub/client.go b/pkg/cyberhub/client.go index 638813c..abc2d50 100644 --- a/pkg/cyberhub/client.go +++ b/pkg/cyberhub/client.go @@ -36,12 +36,9 @@ func newClient(baseURL, apiKey string, timeout time.Duration) *client { } } -func (c *client) exportFingers(ctx context.Context, filter *ExportFilter, draft bool) (types.Fingers, []*types.Alias, error) { +func (c *client) exportFingers(ctx context.Context, filter *ExportFilter) (types.Fingers, []*types.Alias, error) { params := url.Values{} params.Set("with_fingerprint", "true") - if draft { - params.Set("with_draft", "true") - } applyFilterParams(params, filter) endpoint := fmt.Sprintf("%s/fingerprints/export?%s", c.baseURL, params.Encode()) @@ -64,11 +61,8 @@ func (c *client) exportFingers(ctx context.Context, filter *ExportFilter, draft return allFingers, allAliases, nil } -func (c *client) exportPOCs(ctx context.Context, filter *ExportFilter, draft bool) ([]pocResponse, error) { +func (c *client) exportPOCs(ctx context.Context, filter *ExportFilter) ([]pocResponse, error) { params := url.Values{} - if draft { - params.Set("with_draft", "true") - } applyFilterParams(params, filter) applyDefaultPOCStatus(params) @@ -120,6 +114,14 @@ func applyFilterParams(params url.Values, filter *ExportFilter) { params.Set("review_status", filter.ReviewStatus) } + // Draft is orthogonal to Statuses / ReviewStatus: filter fields choose + // the rows, with_draft chooses the column read off each row. Only emit + // the query param when the caller explicitly asked for drafts so the + // default backend semantics (RawContent) is preserved. + if filter.Draft { + params.Set("with_draft", "true") + } + if filter.CreatedAfter != nil { params.Set("created_after", filter.CreatedAfter.Format(time.RFC3339)) } diff --git a/pkg/cyberhub/client_test.go b/pkg/cyberhub/client_test.go index 018ebf6..d661bbf 100644 --- a/pkg/cyberhub/client_test.go +++ b/pkg/cyberhub/client_test.go @@ -1,8 +1,13 @@ package cyberhub import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" "net/url" "testing" + "time" ) func TestApplyFilterParams_DedupTags(t *testing.T) { @@ -111,6 +116,65 @@ func TestApplyFilterParams_ReviewStatus(t *testing.T) { } } +func TestApplyFilterParams_Draft(t *testing.T) { + t.Run("default false omits with_draft", func(t *testing.T) { + params := url.Values{} + applyFilterParams(params, &ExportFilter{}) + + if _, ok := params["with_draft"]; ok { + t.Fatalf("expected no with_draft param when Draft=false, got %v", params["with_draft"]) + } + }) + + t.Run("Draft=true transmits with_draft=true", func(t *testing.T) { + params := url.Values{} + applyFilterParams(params, &ExportFilter{Draft: true}) + + if got := params.Get("with_draft"); got != "true" { + t.Fatalf("expected with_draft=true, got %q", got) + } + }) + + t.Run("Draft orthogonal to Statuses/ReviewStatus", func(t *testing.T) { + params := url.Values{} + filter := &ExportFilter{ + Statuses: []string{"pending", "draft"}, + ReviewStatus: "pending", + Draft: true, + } + applyFilterParams(params, filter) + + if got := params.Get("with_draft"); got != "true" { + t.Fatalf("expected with_draft=true, got %q", got) + } + if got := params.Get("review_status"); got != "pending" { + t.Fatalf("expected review_status=pending, got %q", got) + } + if got := params["statuses"]; len(got) != 2 { + t.Fatalf("expected 2 statuses, got %v", got) + } + }) + + t.Run("nil filter does not panic and omits with_draft", func(t *testing.T) { + params := url.Values{} + applyFilterParams(params, nil) + + if _, ok := params["with_draft"]; ok { + t.Fatalf("expected no with_draft param for nil filter, got %v", params["with_draft"]) + } + }) +} + +func TestWithDraftBuilder(t *testing.T) { + filter := NewExportFilter().WithDraft(true) + if !filter.Draft { + t.Fatalf("expected Draft=true after WithDraft(true)") + } + if got := filter.WithDraft(false); got.Draft { + t.Fatalf("expected WithDraft(false) to clear the flag") + } +} + func TestApplyFilterParams_Limit(t *testing.T) { params := url.Values{} filter := &ExportFilter{Limit: 10} @@ -152,3 +216,102 @@ func TestApplyDefaultPOCStatus_WithReviewStatus(t *testing.T) { t.Fatalf("expected no default status when review_status set, got %q", params.Get("status")) } } + +// TestProvider_DraftBehavior verifies that ExportFilter.Draft reaches the +// server as ?with_draft=true via both Provider.Fingers and Provider.POCs. +func TestProvider_DraftBehavior(t *testing.T) { + var captured url.Values + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + captured = r.URL.Query() + var data interface{} + switch { + case r.URL.Path == "/api/v1/fingerprints/export": + data = fingerprintListResponse{} + case r.URL.Path == "/api/v1/pocs/export": + data = pocListResponse{} + } + resp := apiResponse{Code: 0, Message: "ok", Data: data} + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(resp) + })) + defer server.Close() + + ctx := context.Background() + + t.Run("fingers without WithDraft omits with_draft", func(t *testing.T) { + captured = nil + p := NewProvider(server.URL, "test-key").WithTimeout(5 * time.Second) + if _, _, err := p.Fingers(ctx); err != nil { + t.Fatalf("Fingers failed: %v", err) + } + if _, ok := captured["with_draft"]; ok { + t.Fatalf("expected no with_draft param by default, got %v", captured["with_draft"]) + } + }) + + t.Run("fingers with WithDraft(true) sends with_draft=true", func(t *testing.T) { + captured = nil + filter := NewExportFilter().WithDraft(true).WithReviewStatus("pending") + p := NewProvider(server.URL, "test-key").WithFilter(filter).WithTimeout(5 * time.Second) + if _, _, err := p.Fingers(ctx); err != nil { + t.Fatalf("Fingers failed: %v", err) + } + if got := captured.Get("with_draft"); got != "true" { + t.Fatalf("expected with_draft=true, got %q", got) + } + if got := captured.Get("review_status"); got != "pending" { + t.Fatalf("expected review_status=pending, got %q", got) + } + }) + + t.Run("pocs without WithDraft omits with_draft and keeps default active", func(t *testing.T) { + captured = nil + p := NewProvider(server.URL, "test-key").WithTimeout(5 * time.Second) + if _, err := p.POCs(ctx); err != nil { + t.Fatalf("POCs failed: %v", err) + } + if _, ok := captured["with_draft"]; ok { + t.Fatalf("expected no with_draft param by default, got %v", captured["with_draft"]) + } + if got := captured.Get("status"); got != "active" { + t.Fatalf("expected default status=active, got %q", got) + } + }) + + t.Run("pocs with WithDraft(true) alone keeps default active but adds with_draft=true", func(t *testing.T) { + captured = nil + filter := NewExportFilter().WithDraft(true) + p := NewProvider(server.URL, "test-key").WithFilter(filter).WithTimeout(5 * time.Second) + if _, err := p.POCs(ctx); err != nil { + t.Fatalf("POCs failed: %v", err) + } + if got := captured.Get("with_draft"); got != "true" { + t.Fatalf("expected with_draft=true, got %q", got) + } + // Draft alone does not change the row-filter; caller must add + // WithStatuses / WithReviewStatus to pull pending rows. + if got := captured.Get("status"); got != "active" { + t.Fatalf("expected default status=active when Draft=true without status filter, got %q", got) + } + }) + + t.Run("pocs with WithReviewStatus + WithDraft pulls pending drafts", func(t *testing.T) { + captured = nil + filter := NewExportFilter().WithReviewStatus("pending").WithDraft(true) + p := NewProvider(server.URL, "test-key").WithFilter(filter).WithTimeout(5 * time.Second) + if _, err := p.POCs(ctx); err != nil { + t.Fatalf("POCs failed: %v", err) + } + if got := captured.Get("with_draft"); got != "true" { + t.Fatalf("expected with_draft=true, got %q", got) + } + if got := captured.Get("review_status"); got != "pending" { + t.Fatalf("expected review_status=pending, got %q", got) + } + // review_status suppresses the default active. + if got := captured.Get("status"); got != "" { + t.Fatalf("expected status= empty, got %q", got) + } + }) +} diff --git a/pkg/cyberhub/provider.go b/pkg/cyberhub/provider.go index 8cdecd4..ab1bbde 100644 --- a/pkg/cyberhub/provider.go +++ b/pkg/cyberhub/provider.go @@ -14,7 +14,6 @@ type Provider struct { apiKey string timeout time.Duration filter *ExportFilter - draft bool once sync.Once cli *client @@ -41,12 +40,6 @@ func (p *Provider) WithTimeout(d time.Duration) *Provider { return p } -// WithDraft 启用草稿模式：已审核的指纹直接使用，未审核的使用待审核版本，单次 API 调用加载 -func (p *Provider) WithDraft() *Provider { - p.draft = true - return p -} - func (p *Provider) client() *client { p.once.Do(func() { p.cli = newClient(p.url, p.apiKey, p.timeout) @@ -56,12 +49,12 @@ func (p *Provider) client() *client { // Fingers 导出指纹与别名数据 func (p *Provider) Fingers(ctx context.Context) (types.Fingers, []*types.Alias, error) { - return p.client().exportFingers(ctx, p.filter, p.draft) + return p.client().exportFingers(ctx, p.filter) } // POCs 导出 POC 模板数据 func (p *Provider) POCs(ctx context.Context) ([]*types.Template, error) { - responses, err := p.client().exportPOCs(ctx, p.filter, p.draft) + responses, err := p.client().exportPOCs(ctx, p.filter) if err != nil { return nil, err } diff --git a/pkg/types/export_filter.go b/pkg/types/export_filter.go index 920e57c..c01fc42 100644 --- a/pkg/types/export_filter.go +++ b/pkg/types/export_filter.go @@ -12,6 +12,15 @@ type ExportFilter struct { Statuses []string ReviewStatus string + // Draft controls whether unreviewed-draft content is returned in place of + // the approved RawContent. When true the backend sends back RawContentDraft + // for rows that have one (and the approved RawContent otherwise), which is + // the only way to actually read brand-new pending entries or pending edits. + // + // Draft is orthogonal to Statuses / ReviewStatus: filter fields decide + // which rows are returned, Draft decides which column is read. + Draft bool + CreatedAfter *time.Time CreatedBefore *time.Time UpdatedAfter *time.Time @@ -83,3 +92,13 @@ func (f *ExportFilter) WithReviewStatus(status string) *ExportFilter { f.ReviewStatus = status return f } + +// WithDraft toggles draft mode: when true the backend returns RawContentDraft +// for rows that have a pending-review draft (and falls back to RawContent +// otherwise). Combine with WithStatuses("pending"|"draft") or +// WithReviewStatus("pending") to actually pull pending rows — the filter +// fields decide which rows come back, WithDraft decides which column is read. +func (f *ExportFilter) WithDraft(draft bool) *ExportFilter { + f.Draft = draft + return f +}