diff --git a/cmd/audit_logs.go b/cmd/audit_logs.go index 4149500..aa6bd85 100644 --- a/cmd/audit_logs.go +++ b/cmd/audit_logs.go @@ -3,6 +3,7 @@ package cmd import ( "context" "fmt" + "net/http" "strconv" "strings" "time" @@ -17,6 +18,7 @@ import ( type AuditLogsService interface { ListAutoPaging(ctx context.Context, query kernel.AuditLogListParams, opts ...option.RequestOption) *pagination.PageTokenPaginationAutoPager[kernel.AuditLogEntry] + ExportChunk(ctx context.Context, query kernel.AuditLogExportChunkParams, opts ...option.RequestOption) (*http.Response, error) } type AuditLogsCmd struct { diff --git a/cmd/audit_logs_download.go b/cmd/audit_logs_download.go new file mode 100644 index 0000000..a49954e --- /dev/null +++ b/cmd/audit_logs_download.go @@ -0,0 +1,329 @@ +package cmd + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + "io" + "net/http" + "os" + "path/filepath" + "strconv" + "time" + + "github.com/kernel/cli/pkg/util" + "github.com/kernel/kernel-go-sdk" + "github.com/kernel/kernel-go-sdk/option" + "github.com/pterm/pterm" + "github.com/spf13/cobra" +) + +type AuditLogsDownloadInput struct { + Start string + End string + Search string + Method string + ExcludeMethod string + IncludeGet bool + Service string + AuthStrategy string + UserIDs []string + To string + Force bool +} + +const ( + auditLogsDownloadMaxRange = 30 * 24 * time.Hour + auditLogsChunkAttempts = 7 + auditLogsChunkMaxRetryDelay = 8 * time.Second +) + +var auditLogsChunkRetryBaseDelay = time.Second + +func (c AuditLogsCmd) Download(ctx context.Context, in AuditLogsDownloadInput) error { + params, err := buildAuditLogsDownloadParams(in) + if err != nil { + return err + } + + outPath := in.To + if outPath == "" { + outPath = defaultAuditLogsDownloadPath(params.Start, params.End) + } + partialPath := outPath + ".partial" + out, err := openAuditLogsDownloadOutput(partialPath, outPath, in.Force) + if err != nil { + return err + } + cleanupPartial := true + defer func() { + if out != nil { + out.Close() + } + if cleanupPartial { + os.Remove(partialPath) + } + }() + + var cursor string + var bytesWritten, rows int64 + chunks := 0 + for { + if cursor != "" { + params.Cursor = kernel.String(cursor) + } + body, header, err := c.fetchAuditLogsChunk(ctx, params) + if err != nil { + return err + } + chunkRows, nextCursor, hasMore, err := parseAuditLogsChunkHeaders(header, cursor) + if err != nil { + return err + } + if _, err := out.Write(body); err != nil { + return fmt.Errorf("write %s: %w", partialPath, err) + } + cursor = nextCursor + bytesWritten += int64(len(body)) + chunks++ + rows += chunkRows + pterm.Info.Printf("Chunk %d: %d rows (%d total, %s)\n", chunks, chunkRows, rows, util.FormatBytes(bytesWritten)) + if !hasMore { + break + } + } + + if err := out.Sync(); err != nil { + return fmt.Errorf("sync %s: %w", partialPath, err) + } + closeErr := out.Close() + out = nil + if closeErr != nil { + return fmt.Errorf("close %s: %w", partialPath, closeErr) + } + cleanupPartial = false + if err := commitAuditLogsDownloadOutput(partialPath, outPath, in.Force); err != nil { + return fmt.Errorf("%w; completed download remains at %s", err, partialPath) + } + pterm.Success.Printf("Downloaded %d rows (%s) to %s\n", rows, util.FormatBytes(bytesWritten), outPath) + return nil +} + +func (c AuditLogsCmd) fetchAuditLogsChunk(ctx context.Context, params kernel.AuditLogExportChunkParams) ([]byte, http.Header, error) { + for attempt := 1; ; attempt++ { + body, header, err := c.fetchAuditLogsChunkOnce(ctx, params) + if err == nil || attempt == auditLogsChunkAttempts || !retryableAuditLogsChunkError(err) { + return body, header, err + } + delay := min(auditLogsChunkRetryBaseDelay<<(attempt-1), auditLogsChunkMaxRetryDelay) + pterm.Warning.Printf("Chunk download failed (%s); retrying in %s\n", err, delay) + select { + case <-ctx.Done(): + return nil, nil, ctx.Err() + case <-time.After(delay): + } + } +} + +// retryableAuditLogsChunkError reports whether a chunk fetch failure is +// transient. API errors are retried only for 429s and 5xx; everything else +// (network errors, truncated bodies, checksum mismatches) is retried unless +// the context was cancelled. +func retryableAuditLogsChunkError(err error) bool { + var apiErr *kernel.Error + if errors.As(err, &apiErr) { + return apiErr.StatusCode == http.StatusTooManyRequests || apiErr.StatusCode >= 500 + } + return !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) +} + +func (c AuditLogsCmd) fetchAuditLogsChunkOnce(ctx context.Context, params kernel.AuditLogExportChunkParams) ([]byte, http.Header, error) { + res, err := c.auditLogs.ExportChunk(ctx, params, option.WithMaxRetries(0)) + if err != nil { + return nil, nil, util.CleanedUpSdkError{Err: err} + } + defer res.Body.Close() + + body, err := io.ReadAll(res.Body) + if err != nil { + return nil, nil, fmt.Errorf("read chunk body: %w", err) + } + want := res.Header.Get("X-Content-Sha256") + if want == "" { + return nil, nil, fmt.Errorf("response missing X-Content-Sha256 header") + } + sum := sha256.Sum256(body) + if got := hex.EncodeToString(sum[:]); got != want { + return nil, nil, fmt.Errorf("chunk checksum mismatch (got %s, want %s)", got, want) + } + return body, res.Header, nil +} + +func parseAuditLogsChunkHeaders(header http.Header, currentCursor string) (int64, string, bool, error) { + hasMore, err := strconv.ParseBool(header.Get("X-Has-More")) + if err != nil { + return 0, "", false, fmt.Errorf("response missing or invalid X-Has-More header") + } + rows, err := strconv.ParseInt(header.Get("X-Row-Count"), 10, 64) + if err != nil || rows < 0 { + return 0, "", false, fmt.Errorf("response missing or invalid X-Row-Count header") + } + nextCursor := header.Get("X-Next-Cursor") + if hasMore && (nextCursor == "" || nextCursor == currentCursor) { + return 0, "", false, fmt.Errorf("response has invalid X-Next-Cursor header") + } + if !hasMore && nextCursor != "" { + return 0, "", false, fmt.Errorf("response returned a cursor after the final chunk") + } + return rows, nextCursor, hasMore, nil +} + +func buildAuditLogsDownloadParams(in AuditLogsDownloadInput) (kernel.AuditLogExportChunkParams, error) { + var params kernel.AuditLogExportChunkParams + if in.Start == "" || in.End == "" { + return params, fmt.Errorf("--start and --end are required") + } + start, err := parseAuditLogTime(in.Start) + if err != nil { + return params, fmt.Errorf("--start: %w", err) + } + end, err := parseAuditLogTime(in.End) + if err != nil { + return params, fmt.Errorf("--end: %w", err) + } + if !start.Before(end) { + return params, fmt.Errorf("--start must be before --end") + } + if end.Sub(start) > auditLogsDownloadMaxRange { + return params, fmt.Errorf("the API allows at most 30 days per download") + } + + params.Start = start + params.End = end + params.Format = kernel.AuditLogExportChunkParamsFormatJSONLGz + if in.Search != "" { + params.Search = kernel.String(in.Search) + } + if in.Method != "" { + params.Method = kernel.String(in.Method) + } + params.ExcludeMethod = auditLogExcludeMethods(in.Method, in.ExcludeMethod, in.IncludeGet) + if in.Service != "" { + params.Service = kernel.String(in.Service) + } + if in.AuthStrategy != "" { + params.AuthStrategy = kernel.String(in.AuthStrategy) + } + if len(in.UserIDs) > 0 { + params.SearchUserID = in.UserIDs + } + return params, nil +} + +func defaultAuditLogsDownloadPath(start, end time.Time) string { + const stamp = "20060102" + return fmt.Sprintf("audit-logs-%s-%s.jsonl.gz", start.UTC().Format(stamp), end.UTC().Format(stamp)) +} + +// checkAuditLogsDownloadTarget rejects paths that are not replaceable +// regular files, and existing files unless --force was passed. +func checkAuditLogsDownloadTarget(path string, force bool) error { + if info, err := os.Lstat(path); err == nil { + if !info.Mode().IsRegular() { + return fmt.Errorf("output %s is not a regular file", path) + } + if !force { + return fmt.Errorf("output %s already exists; pass --force to overwrite", path) + } + } else if !os.IsNotExist(err) { + return fmt.Errorf("inspect %s: %w", path, err) + } + return nil +} + +func openAuditLogsDownloadOutput(partialPath, outPath string, force bool) (*os.File, error) { + if err := checkAuditLogsDownloadTarget(outPath, force); err != nil { + return nil, err + } + if info, err := os.Lstat(partialPath); err == nil && !info.Mode().IsRegular() { + return nil, fmt.Errorf("partial file %s is not a regular file", partialPath) + } else if err != nil && !os.IsNotExist(err) { + return nil, fmt.Errorf("inspect %s: %w", partialPath, err) + } + if err := os.MkdirAll(filepath.Dir(partialPath), 0o700); err != nil { + return nil, fmt.Errorf("create output directory: %w", err) + } + out, err := os.OpenFile(partialPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o600) + if err != nil { + return nil, fmt.Errorf("open %s: %w", partialPath, err) + } + // OpenFile's mode only applies on creation; a reused partial keeps its + // previous permissions unless tightened explicitly. + if err := out.Chmod(0o600); err != nil { + out.Close() + return nil, fmt.Errorf("secure %s: %w", partialPath, err) + } + return out, nil +} + +func commitAuditLogsDownloadOutput(partialPath, outPath string, force bool) error { + if err := checkAuditLogsDownloadTarget(outPath, force); err != nil { + return err + } + if err := os.Rename(partialPath, outPath); err != nil { + return fmt.Errorf("finalize %s: %w", outPath, err) + } + return nil +} + +func runAuditLogsDownload(cmd *cobra.Command, args []string) error { + c := getAuditLogsHandler(cmd) + start, _ := cmd.Flags().GetString("start") + end, _ := cmd.Flags().GetString("end") + search, _ := cmd.Flags().GetString("search") + method, _ := cmd.Flags().GetString("method") + excludeMethod, _ := cmd.Flags().GetString("exclude-method") + includeGet, _ := cmd.Flags().GetBool("include-get") + service, _ := cmd.Flags().GetString("service") + authStrategy, _ := cmd.Flags().GetString("auth-strategy") + userIDs, _ := cmd.Flags().GetStringArray("user-id") + to, _ := cmd.Flags().GetString("to") + force, _ := cmd.Flags().GetBool("force") + + return c.Download(cmd.Context(), AuditLogsDownloadInput{ + Start: start, End: end, Search: search, Method: method, + ExcludeMethod: excludeMethod, IncludeGet: includeGet, Service: service, + AuthStrategy: authStrategy, UserIDs: userIDs, To: to, Force: force, + }) +} + +var auditLogsDownloadCmd = &cobra.Command{ + Use: "download", + Short: "Download audit logs as gzip-compressed JSONL", + Long: "Download audit logs as gzip-compressed JSONL in verified chunks. The time range is [start, end).\n\n" + + "The API allows at most 30 days per download.\n\n" + + "GET requests are excluded by default; pass --include-get to include them.\n\n" + + "The output file is published only after every chunk is downloaded.", + Example: "download --start 2026-06-01 --end 2026-07-01 --to audit-june.jsonl.gz", + Args: cobra.NoArgs, + RunE: runAuditLogsDownload, +} + +func init() { + auditLogsDownloadCmd.Flags().String("start", "", "Start of the export window, RFC3339 or YYYY-MM-DD (required)") + auditLogsDownloadCmd.Flags().String("end", "", "Exclusive end of the export window, RFC3339 or YYYY-MM-DD (required)") + auditLogsDownloadCmd.Flags().String("search", "", "Free-text search") + auditLogsDownloadCmd.Flags().String("method", "", "Filter by HTTP method") + auditLogsDownloadCmd.Flags().String("exclude-method", "", "Exclude an HTTP method") + auditLogsDownloadCmd.Flags().Bool("include-get", false, "Include GET requests, which are excluded by default") + auditLogsDownloadCmd.Flags().String("service", "", "Filter by service") + auditLogsDownloadCmd.Flags().String("auth-strategy", "", "Filter by authentication strategy") + auditLogsDownloadCmd.Flags().StringArray("user-id", nil, "Filter by user ID (repeatable)") + auditLogsDownloadCmd.Flags().String("to", "", "Output .jsonl.gz file path") + auditLogsDownloadCmd.Flags().Bool("force", false, "Overwrite the output file") + _ = auditLogsDownloadCmd.MarkFlagRequired("start") + _ = auditLogsDownloadCmd.MarkFlagRequired("end") + auditLogsCmd.AddCommand(auditLogsDownloadCmd) +} diff --git a/cmd/audit_logs_download_test.go b/cmd/audit_logs_download_test.go new file mode 100644 index 0000000..9df3d67 --- /dev/null +++ b/cmd/audit_logs_download_test.go @@ -0,0 +1,373 @@ +package cmd + +import ( + "bytes" + "compress/gzip" + "context" + "crypto/sha256" + "encoding/hex" + "errors" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strconv" + "sync/atomic" + "testing" + "time" + + "github.com/kernel/kernel-go-sdk" + "github.com/kernel/kernel-go-sdk/option" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type auditLogTestChunk struct { + body []byte + rows int + hasMore bool + nextCursor string + checksum string +} + +func auditLogChunkResponse(chunk auditLogTestChunk) *http.Response { + checksum := chunk.checksum + if checksum == "" { + sum := sha256.Sum256(chunk.body) + checksum = hex.EncodeToString(sum[:]) + } + header := http.Header{} + header.Set("X-Content-Sha256", checksum) + header.Set("X-Row-Count", strconv.Itoa(chunk.rows)) + header.Set("X-Has-More", strconv.FormatBool(chunk.hasMore)) + if chunk.nextCursor != "" { + header.Set("X-Next-Cursor", chunk.nextCursor) + } + return &http.Response{StatusCode: http.StatusOK, Header: header, Body: io.NopCloser(bytes.NewReader(chunk.body))} +} + +func auditLogChunkService(t *testing.T, responses ...func() (*http.Response, error)) (*FakeAuditLogsService, *[]string) { + t.Helper() + call := 0 + cursors := make([]string, 0, len(responses)) + service := &FakeAuditLogsService{ + ExportChunkFunc: func(ctx context.Context, query kernel.AuditLogExportChunkParams, opts ...option.RequestOption) (*http.Response, error) { + cursors = append(cursors, query.Cursor.Value) + require.Less(t, call, len(responses)) + response, err := responses[call]() + call++ + return response, err + }, + } + return service, &cursors +} + +func auditLogGzip(t *testing.T, body string) []byte { + t.Helper() + var out bytes.Buffer + writer := gzip.NewWriter(&out) + _, err := writer.Write([]byte(body)) + require.NoError(t, err) + require.NoError(t, writer.Close()) + return out.Bytes() +} + +func readAuditLogGzip(t *testing.T, path string) string { + t.Helper() + data, err := os.ReadFile(path) + require.NoError(t, err) + reader, err := gzip.NewReader(bytes.NewReader(data)) + require.NoError(t, err) + body, err := io.ReadAll(reader) + require.NoError(t, err) + return string(body) +} + +func auditLogsDownloadInput(path string) AuditLogsDownloadInput { + return AuditLogsDownloadInput{Start: "2026-06-01", End: "2026-06-28", To: path} +} + +func disableAuditLogsRetryDelay(t *testing.T) { + t.Helper() + saved := auditLogsChunkRetryBaseDelay + auditLogsChunkRetryBaseDelay = 0 + t.Cleanup(func() { auditLogsChunkRetryBaseDelay = saved }) +} + +func TestAuditLogsDownloadWritesAllChunks(t *testing.T) { + capturePtermOutput(t) + outPath := filepath.Join(t.TempDir(), "audit.jsonl.gz") + first := auditLogGzip(t, "{\"n\":1}\n") + second := auditLogGzip(t, "{\"n\":2}\n") + service, cursors := auditLogChunkService(t, + func() (*http.Response, error) { + return auditLogChunkResponse(auditLogTestChunk{body: first, rows: 1, hasMore: true, nextCursor: "next"}), nil + }, + func() (*http.Response, error) { + return auditLogChunkResponse(auditLogTestChunk{body: second, rows: 1}), nil + }, + ) + + err := (AuditLogsCmd{auditLogs: service}).Download(context.Background(), auditLogsDownloadInput(outPath)) + require.NoError(t, err) + assert.Equal(t, []string{"", "next"}, *cursors) + assert.Equal(t, "{\"n\":1}\n{\"n\":2}\n", readAuditLogGzip(t, outPath)) + _, err = os.Stat(outPath + ".partial") + assert.True(t, os.IsNotExist(err)) +} + +func TestAuditLogsDownloadExcludeMethodStacksWithDefaultGetExclusion(t *testing.T) { + capturePtermOutput(t) + outPath := filepath.Join(t.TempDir(), "audit.jsonl.gz") + body := auditLogGzip(t, "{\"method\":\"DELETE\",\"n\":2}\n") + service := &FakeAuditLogsService{ + ExportChunkFunc: func(ctx context.Context, query kernel.AuditLogExportChunkParams, opts ...option.RequestOption) (*http.Response, error) { + assert.Equal(t, []string{"GET", "post"}, query.ExcludeMethod) + return auditLogChunkResponse(auditLogTestChunk{body: body, rows: 1}), nil + }, + } + in := auditLogsDownloadInput(outPath) + in.ExcludeMethod = "post" + + require.NoError(t, (AuditLogsCmd{auditLogs: service}).Download(context.Background(), in)) + assert.Equal(t, "{\"method\":\"DELETE\",\"n\":2}\n", readAuditLogGzip(t, outPath)) +} + +func TestAuditLogsDownloadRemovesPartialAfterTransferFailure(t *testing.T) { + capturePtermOutput(t) + disableAuditLogsRetryDelay(t) + outPath := filepath.Join(t.TempDir(), "audit.jsonl.gz") + first := auditLogGzip(t, "{\"n\":1}\n") + fail := func() (*http.Response, error) { return nil, errors.New("network error") } + service, _ := auditLogChunkService(t, + func() (*http.Response, error) { + return auditLogChunkResponse(auditLogTestChunk{body: first, rows: 1, hasMore: true, nextCursor: "next"}), nil + }, + fail, fail, fail, fail, fail, fail, fail, + ) + require.Error(t, (AuditLogsCmd{auditLogs: service}).Download(context.Background(), auditLogsDownloadInput(outPath))) + _, err := os.Stat(outPath) + assert.True(t, os.IsNotExist(err)) + _, err = os.Stat(outPath + ".partial") + assert.True(t, os.IsNotExist(err)) +} + +func TestAuditLogsDownloadOverwritesStalePartial(t *testing.T) { + capturePtermOutput(t) + outPath := filepath.Join(t.TempDir(), "audit.jsonl.gz") + partialPath := outPath + ".partial" + require.NoError(t, os.WriteFile(partialPath, []byte("stale"), 0o644)) + + chunk := auditLogGzip(t, "{\"n\":1}\n") + service, cursors := auditLogChunkService(t, func() (*http.Response, error) { + return auditLogChunkResponse(auditLogTestChunk{body: chunk, rows: 1}), nil + }) + + require.NoError(t, (AuditLogsCmd{auditLogs: service}).Download(context.Background(), auditLogsDownloadInput(outPath))) + assert.Equal(t, []string{""}, *cursors) + assert.Equal(t, "{\"n\":1}\n", readAuditLogGzip(t, outPath)) + info, err := os.Stat(outPath) + require.NoError(t, err) + assert.Equal(t, os.FileMode(0o600), info.Mode().Perm()) + _, err = os.Stat(partialPath) + assert.True(t, os.IsNotExist(err)) +} + +func TestAuditLogsDownloadRetriesFailedChunk(t *testing.T) { + capturePtermOutput(t) + disableAuditLogsRetryDelay(t) + outPath := filepath.Join(t.TempDir(), "audit.jsonl.gz") + first := auditLogGzip(t, "{\"n\":1}\n") + second := auditLogGzip(t, "{\"n\":2}\n") + service, cursors := auditLogChunkService(t, + func() (*http.Response, error) { + return auditLogChunkResponse(auditLogTestChunk{body: first, rows: 1, hasMore: true, nextCursor: "next"}), nil + }, + func() (*http.Response, error) { return nil, errors.New("network error") }, + func() (*http.Response, error) { + return auditLogChunkResponse(auditLogTestChunk{body: second, rows: 1}), nil + }, + ) + + require.NoError(t, (AuditLogsCmd{auditLogs: service}).Download(context.Background(), auditLogsDownloadInput(outPath))) + assert.Equal(t, []string{"", "next", "next"}, *cursors) + assert.Equal(t, "{\"n\":1}\n{\"n\":2}\n", readAuditLogGzip(t, outPath)) +} + +func TestAuditLogsDownloadRecoversFromChecksumMismatch(t *testing.T) { + capturePtermOutput(t) + disableAuditLogsRetryDelay(t) + outPath := filepath.Join(t.TempDir(), "audit.jsonl.gz") + chunk := auditLogGzip(t, "{\"n\":1}\n") + service, cursors := auditLogChunkService(t, + func() (*http.Response, error) { + return auditLogChunkResponse(auditLogTestChunk{body: []byte("bad"), rows: 1, checksum: "wrong"}), nil + }, + func() (*http.Response, error) { + return auditLogChunkResponse(auditLogTestChunk{body: chunk, rows: 1}), nil + }, + ) + + require.NoError(t, (AuditLogsCmd{auditLogs: service}).Download(context.Background(), auditLogsDownloadInput(outPath))) + assert.Equal(t, []string{"", ""}, *cursors) + assert.Equal(t, "{\"n\":1}\n", readAuditLogGzip(t, outPath)) +} + +func TestAuditLogsDownloadDisablesSDKRetries(t *testing.T) { + capturePtermOutput(t) + disableAuditLogsRetryDelay(t) + var requests atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requests.Add(1) + w.WriteHeader(http.StatusInternalServerError) + })) + defer server.Close() + client := kernel.NewClient(option.WithBaseURL(server.URL), option.WithAPIKey("test")) + outPath := filepath.Join(t.TempDir(), "audit.jsonl.gz") + + err := (AuditLogsCmd{auditLogs: &client.AuditLogs}).Download(context.Background(), auditLogsDownloadInput(outPath)) + require.Error(t, err) + assert.Equal(t, int64(auditLogsChunkAttempts), requests.Load()) +} + +func TestAuditLogsDownloadDoesNotRetryClientErrors(t *testing.T) { + capturePtermOutput(t) + disableAuditLogsRetryDelay(t) + outPath := filepath.Join(t.TempDir(), "audit.jsonl.gz") + service, _ := auditLogChunkService(t, func() (*http.Response, error) { + return nil, &kernel.Error{StatusCode: http.StatusBadRequest} + }) + + err := (AuditLogsCmd{auditLogs: service}).Download(context.Background(), auditLogsDownloadInput(outPath)) + require.Error(t, err) +} + +func TestDefaultAuditLogsDownloadPath(t *testing.T) { + start := time.Date(2026, 6, 1, 0, 0, 0, 0, time.UTC) + end := time.Date(2026, 6, 28, 0, 0, 0, 0, time.UTC) + path := defaultAuditLogsDownloadPath(start, end) + + assert.Equal(t, "audit-logs-20260601-20260628.jsonl.gz", path) +} + +func TestAuditLogsDownloadRejectsBadChunkBeforeWriting(t *testing.T) { + capturePtermOutput(t) + disableAuditLogsRetryDelay(t) + outPath := filepath.Join(t.TempDir(), "audit.jsonl.gz") + badChunk := func() (*http.Response, error) { + return auditLogChunkResponse(auditLogTestChunk{body: []byte("bad"), rows: 1, checksum: "wrong"}), nil + } + service, _ := auditLogChunkService(t, badChunk, badChunk, badChunk, badChunk, badChunk, badChunk, badChunk) + + err := (AuditLogsCmd{auditLogs: service}).Download(context.Background(), auditLogsDownloadInput(outPath)) + require.ErrorContains(t, err, "checksum mismatch") + _, statErr := os.Stat(outPath) + assert.True(t, os.IsNotExist(statErr)) + _, statErr = os.Stat(outPath + ".partial") + assert.True(t, os.IsNotExist(statErr)) +} + +func TestAuditLogsDownloadForceOverwrites(t *testing.T) { + capturePtermOutput(t) + outPath := filepath.Join(t.TempDir(), "audit.jsonl.gz") + require.NoError(t, os.WriteFile(outPath, []byte("old"), 0o600)) + chunk := auditLogGzip(t, "{\"n\":1}\n") + service, cursors := auditLogChunkService(t, func() (*http.Response, error) { + return auditLogChunkResponse(auditLogTestChunk{body: chunk, rows: 1}), nil + }) + in := auditLogsDownloadInput(outPath) + in.Force = true + + require.NoError(t, (AuditLogsCmd{auditLogs: service}).Download(context.Background(), in)) + assert.Equal(t, []string{""}, *cursors) + assert.Equal(t, "{\"n\":1}\n", readAuditLogGzip(t, outPath)) +} + +func TestAuditLogsDownloadPreservesCompletedPartialOnFinalizeFailure(t *testing.T) { + capturePtermOutput(t) + dir := t.TempDir() + outPath := filepath.Join(dir, "audit.jsonl.gz") + chunk := auditLogGzip(t, "{\"n\":1}\n") + service, _ := auditLogChunkService(t, func() (*http.Response, error) { + require.NoError(t, os.Mkdir(outPath, 0o700)) + return auditLogChunkResponse(auditLogTestChunk{body: chunk, rows: 1}), nil + }) + + err := (AuditLogsCmd{auditLogs: service}).Download(context.Background(), auditLogsDownloadInput(outPath)) + require.ErrorContains(t, err, "completed download remains") + assert.Equal(t, "{\"n\":1}\n", readAuditLogGzip(t, outPath+".partial")) +} + +func TestCommitAuditLogsDownloadOutputPreservesExistingFileOnFailure(t *testing.T) { + outPath := filepath.Join(t.TempDir(), "audit.jsonl.gz") + require.NoError(t, os.WriteFile(outPath, []byte("existing"), 0o600)) + + err := commitAuditLogsDownloadOutput(outPath+".partial", outPath, true) + require.Error(t, err) + data, readErr := os.ReadFile(outPath) + require.NoError(t, readErr) + assert.Equal(t, "existing", string(data)) +} + +func TestBuildAuditLogsDownloadParams(t *testing.T) { + in := auditLogsDownloadInput("") + in.Search = "browser" + in.Service = "api" + in.UserIDs = []string{"user_1"} + params, err := buildAuditLogsDownloadParams(in) + require.NoError(t, err) + + assert.Equal(t, time.Date(2026, 6, 28, 0, 0, 0, 0, time.UTC), params.End) + assert.Equal(t, "browser", params.Search.Value) + assert.Equal(t, []string{"GET"}, params.ExcludeMethod) + assert.Equal(t, "api", params.Service.Value) + assert.Equal(t, []string{"user_1"}, params.SearchUserID) +} + +func TestBuildAuditLogsDownloadParamsRejectsInvalidInput(t *testing.T) { + tests := []struct { + name string + in AuditLogsDownloadInput + want string + }{ + {name: "missing bounds", in: AuditLogsDownloadInput{}, want: "--start and --end are required"}, + {name: "reversed bounds", in: AuditLogsDownloadInput{Start: "2026-06-02", End: "2026-06-01"}, want: "--start must be before --end"}, + {name: "range too large", in: AuditLogsDownloadInput{Start: "2026-05-01", End: "2026-06-01"}, want: "at most 30 days"}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + _, err := buildAuditLogsDownloadParams(test.in) + require.ErrorContains(t, err, test.want) + }) + } +} + +func TestParseAuditLogsChunkHeaders(t *testing.T) { + tests := []struct { + name string + header http.Header + current string + want string + }{ + {name: "missing has more", header: http.Header{"X-Row-Count": []string{"1"}}, want: "X-Has-More"}, + {name: "missing row count", header: http.Header{"X-Has-More": []string{"false"}}, want: "X-Row-Count"}, + {name: "missing cursor", header: http.Header{"X-Has-More": []string{"true"}, "X-Row-Count": []string{"1"}}, want: "X-Next-Cursor"}, + {name: "unchanged cursor", current: "next", header: http.Header{"X-Has-More": []string{"true"}, "X-Row-Count": []string{"1"}, "X-Next-Cursor": []string{"next"}}, want: "X-Next-Cursor"}, + {name: "cursor after final chunk", header: http.Header{"X-Has-More": []string{"false"}, "X-Row-Count": []string{"1"}, "X-Next-Cursor": []string{"next"}}, want: "cursor after the final chunk"}, + {name: "negative row count", header: http.Header{"X-Has-More": []string{"false"}, "X-Row-Count": []string{"-1"}}, want: "X-Row-Count"}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + _, _, _, err := parseAuditLogsChunkHeaders(test.header, test.current) + require.ErrorContains(t, err, test.want) + }) + } +} + +func TestAuditLogsDownloadRejectsExistingOutput(t *testing.T) { + outPath := filepath.Join(t.TempDir(), "audit.jsonl.gz") + require.NoError(t, os.WriteFile(outPath, []byte("keep"), 0o600)) + err := (AuditLogsCmd{auditLogs: &FakeAuditLogsService{}}).Download(context.Background(), auditLogsDownloadInput(outPath)) + require.ErrorContains(t, err, "already exists") +} diff --git a/cmd/audit_logs_test.go b/cmd/audit_logs_test.go index 8409393..32a1d47 100644 --- a/cmd/audit_logs_test.go +++ b/cmd/audit_logs_test.go @@ -17,6 +17,7 @@ import ( type FakeAuditLogsService struct { ListAutoPagingFunc func(ctx context.Context, query kernel.AuditLogListParams, opts ...option.RequestOption) *pagination.PageTokenPaginationAutoPager[kernel.AuditLogEntry] + ExportChunkFunc func(ctx context.Context, query kernel.AuditLogExportChunkParams, opts ...option.RequestOption) (*http.Response, error) } func (f *FakeAuditLogsService) ListAutoPaging(ctx context.Context, query kernel.AuditLogListParams, opts ...option.RequestOption) *pagination.PageTokenPaginationAutoPager[kernel.AuditLogEntry] { @@ -26,6 +27,13 @@ func (f *FakeAuditLogsService) ListAutoPaging(ctx context.Context, query kernel. return auditLogPager() } +func (f *FakeAuditLogsService) ExportChunk(ctx context.Context, query kernel.AuditLogExportChunkParams, opts ...option.RequestOption) (*http.Response, error) { + if f.ExportChunkFunc != nil { + return f.ExportChunkFunc(ctx, query, opts...) + } + return nil, errors.New("ExportChunk not implemented") +} + func auditLogPager(entries ...kernel.AuditLogEntry) *pagination.PageTokenPaginationAutoPager[kernel.AuditLogEntry] { page := &pagination.PageTokenPagination[kernel.AuditLogEntry]{Items: entries} page.SetPageConfig(nil, &http.Response{Header: http.Header{}})