From 05c0569dae8668e513986b3283d8d863eb8eb6c4 Mon Sep 17 00:00:00 2001 From: Gaziz Nugmanov Date: Thu, 11 Sep 2025 17:02:51 -0700 Subject: [PATCH] SNOW-2334154: add ip-only/no-ip prefix for output format --- cmd/sanssh/client/client.go | 11 +++ cmd/sanssh/main.go | 34 ++++--- services/util/writer/clean-writer.go | 82 +++++++++++++++++ services/util/writer/iponly-writer.go | 100 +++++++++++++++++++++ services/util/writer/prefix-writer_test.go | 84 +++++++++++++++++ 5 files changed, 297 insertions(+), 14 deletions(-) create mode 100644 services/util/writer/clean-writer.go create mode 100644 services/util/writer/iponly-writer.go diff --git a/cmd/sanssh/client/client.go b/cmd/sanssh/client/client.go index 6e2102ff..14b53a0b 100644 --- a/cmd/sanssh/client/client.go +++ b/cmd/sanssh/client/client.go @@ -75,6 +75,10 @@ type RunState struct { ClientAuthzPolicy rpcauth.AuthzPolicy // PrefixOutput if true will prefix every line of output with '-: ' PrefixOutput bool + // CleanOutputIPOnly if true will output only the remote host/IP per line + CleanOutputIPOnly bool + // CleanOutput if true will strip the first token up to the first space from each output line + CleanOutput bool // BatchSize if non-zero will do the requested operation to the targets but in // N calls to the proxy where N is the target list size divided by BatchSize. BatchSize int @@ -308,6 +312,13 @@ func Run(ctx context.Context, rs RunState) { } makeWriter := func(prefix bool, i int, dest io.Writer) io.Writer { + // Optionally clean original output lines first (IP-only takes precedence) + if rs.CleanOutputIPOnly { + dest = writerUtils.GetIPOnlyWriter(dest) + } else if rs.CleanOutput { + dest = writerUtils.GetCleanOutputWriter(dest) + } + // Then add our prefix (if requested) if prefix { targetName := cmdUtil.StripTimeout(rs.Targets[i]) dest = writerUtils.GetPrefixedWriter( diff --git a/cmd/sanssh/main.go b/cmd/sanssh/main.go index a209174e..3338843c 100644 --- a/cmd/sanssh/main.go +++ b/cmd/sanssh/main.go @@ -76,20 +76,22 @@ var ( If blank a direct connection to the first entry in --targets will be made. If port is blank the default of %d will be used`, proxyEnv, defaultProxyPort)) // Deprecated: --timeout flag is deprecated. Use --idle-timeout or --dial-timeout instead - _ = flag.Duration("timeout", defaultDialTimeout, "DEPRECATED. Please use --idle-timeout or --dial-timeout instead") - dialTimeout = flag.Duration("dial-timeout", defaultDialTimeout, "How long to wait for the connection to be accepted. Timeout specified in --targets or --proxy will take precedence") - idleTimeout = flag.Duration("idle-timeout", defaultIdleTimeout, "Maximum time that a connection is idle. If no messages are received within this timeframe, connection will be terminated") - credSource = flag.String("credential-source", mtlsFlags.Name(), fmt.Sprintf("Method used to obtain mTLS credentials (one of [%s])", strings.Join(mtls.Loaders(), ","))) - outputsDir = flag.String("output-dir", "", "If set defines a directory to emit output/errors from commands. Files will be generated based on target as destination/0 destination/0.error, etc.") - justification = flag.String("justification", "", "If non-empty will add the key '"+rpcauth.ReqJustKey+"' to the outgoing context Metadata to be passed along to the server for possible validation and logging.") - targetsFile = flag.String("targets-file", "", "If set read the targets list line by line (as host[:port]) from the indicated file instead of using --targets (error if both flags are used). A blank port acts the same as --targets") - clientPolicyFlag = flag.String("client-policy", "", "OPA policy for outbound client actions. If empty no policy is applied.") - clientPolicyFile = flag.String("client-policy-file", "", "Path to a file with a client OPA. If empty uses --client-policy") - verbosity = flag.Int("v", -1, "Verbosity level. > 0 indicates more extensive logging") - prefixHeader = flag.Bool("h", false, "If true prefix each line of output with '-: '") - batchSize = flag.Int("batch-size", 0, "If non-zero will perform the proxy->target work in batches of this size (with any remainder done at the end).") - mpa = flag.Bool("mpa", false, "Request multi-party approval for commands. This will create an MPA request, wait for approval, and then execute the command.") - authzDryRun = flag.Bool("authz-dry-run", false, "If true, the client will send a request to the server to check if the user has the permission to run the command. The server will respond with a success or failure message.") + _ = flag.Duration("timeout", defaultDialTimeout, "DEPRECATED. Please use --idle-timeout or --dial-timeout instead") + dialTimeout = flag.Duration("dial-timeout", defaultDialTimeout, "How long to wait for the connection to be accepted. Timeout specified in --targets or --proxy will take precedence") + idleTimeout = flag.Duration("idle-timeout", defaultIdleTimeout, "Maximum time that a connection is idle. If no messages are received within this timeframe, connection will be terminated") + credSource = flag.String("credential-source", mtlsFlags.Name(), fmt.Sprintf("Method used to obtain mTLS credentials (one of [%s])", strings.Join(mtls.Loaders(), ","))) + outputsDir = flag.String("output-dir", "", "If set defines a directory to emit output/errors from commands. Files will be generated based on target as destination/0 destination/0.error, etc.") + justification = flag.String("justification", "", "If non-empty will add the key '"+rpcauth.ReqJustKey+"' to the outgoing context Metadata to be passed along to the server for possible validation and logging.") + targetsFile = flag.String("targets-file", "", "If set read the targets list line by line (as host[:port]) from the indicated file instead of using --targets (error if both flags are used). A blank port acts the same as --targets") + clientPolicyFlag = flag.String("client-policy", "", "OPA policy for outbound client actions. If empty no policy is applied.") + clientPolicyFile = flag.String("client-policy-file", "", "Path to a file with a client OPA. If empty uses --client-policy") + verbosity = flag.Int("v", -1, "Verbosity level. > 0 indicates more extensive logging") + prefixHeader = flag.Bool("h", false, "If true prefix each line of output with '-: '") + cleanOutput = flag.Bool("clean-output", false, "If true, strip the first token up to the first space from each output line") + cleanOutputIPOnly = flag.Bool("clean-output-ip-only", false, "If true, print only the remote host/IP per output line") + batchSize = flag.Int("batch-size", 0, "If non-zero will perform the proxy->target work in batches of this size (with any remainder done at the end).") + mpa = flag.Bool("mpa", false, "Request multi-party approval for commands. This will create an MPA request, wait for approval, and then execute the command.") + authzDryRun = flag.Bool("authz-dry-run", false, "If true, the client will send a request to the server to check if the user has the permission to run the command. The server will respond with a success or failure message.") // targets will be bound to --targets for sending a single request to N nodes. targetsFlag util.StringSliceCommaOrWhitespaceFlag @@ -126,6 +128,8 @@ func init() { subcommands.ImportantFlag("client-policy-file") subcommands.ImportantFlag("mpa") subcommands.ImportantFlag("v") + subcommands.ImportantFlag("clean-output") + subcommands.ImportantFlag("clean-output-ip-only") } func isFlagPassed(name string) bool { @@ -218,6 +222,8 @@ func main() { IdleTimeout: *idleTimeout, ClientAuthzPolicy: clientPolicy, PrefixOutput: *prefixHeader, + CleanOutput: *cleanOutput, + CleanOutputIPOnly: *cleanOutputIPOnly, BatchSize: *batchSize, EnableMPA: *mpa, } diff --git a/services/util/writer/clean-writer.go b/services/util/writer/clean-writer.go new file mode 100644 index 00000000..611d3f5d --- /dev/null +++ b/services/util/writer/clean-writer.go @@ -0,0 +1,82 @@ +/* Copyright (c) 2025 Snowflake Inc. All rights reserved. + + Licensed under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. +*/ + +package writer + +import ( + "bytes" + "io" +) + +// GetCleanOutputWriter returns a writer that, for each line written to it, +// removes everything from the beginning of the line up to and including the +// first space character, then writes the remainder to the underlying writer. +// +// This is useful for transforming lines like: +// +// "123-10.0.0.1:9500: some_output_strings\n" +// +// into: +// +// "some_output_strings\n" +// +// The writer is newline-aware and will correctly handle content split across +// multiple Write() calls. +func GetCleanOutputWriter(dest io.Writer) WrappedWriter { + return &cleanWriter{ + dest: dest, + } +} + +type cleanWriter struct { + dest io.Writer + buf []byte +} + +func (c *cleanWriter) Write(p []byte) (n int, err error) { + if len(p) == 0 { + return 0, nil + } + // Keep track of input length to report as written count regardless of + // transformation performed. + total := len(p) + + // Accumulate into buffer and flush full lines. + c.buf = append(c.buf, p...) + for { + idx := bytes.IndexByte(c.buf, '\n') + if idx == -1 { + break + } + line := c.buf[:idx+1] // include the newline + + // Find the first space and strip everything up to and including it. + out := line + if sp := bytes.IndexByte(line, ' '); sp >= 0 { + out = line[sp+1:] + } + if _, err := c.dest.Write(out); err != nil { + return 0, err + } + // Advance buffer + c.buf = c.buf[idx+1:] + } + return total, nil +} + +func (c *cleanWriter) GetOriginal() io.Writer { + return c.dest +} diff --git a/services/util/writer/iponly-writer.go b/services/util/writer/iponly-writer.go new file mode 100644 index 00000000..f89ec919 --- /dev/null +++ b/services/util/writer/iponly-writer.go @@ -0,0 +1,100 @@ +/* Copyright (c) 2025 Snowflake Inc. All rights reserved. + + Licensed under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. +*/ + +package writer + +import ( + "bytes" + "io" +) + +// GetIPOnlyWriter returns a writer that, for each line written to it, +// extracts the remote host IP (or host) from a header of the form +// "-:: ..." and writes only the IP/host followed by a newline. +// +// Examples: +// +// "123-10.0.0.1:9500: some_output_strings\n" -> "10.0.0.1\n" +// "77-[2001:db8::1]:50042: foo\n" -> "2001:db8::1\n" +// "9-localhost:50042: bar\n" -> "localhost\n" +// +// The writer is newline-aware and will correctly handle content split across multiple writes. +func GetIPOnlyWriter(dest io.Writer) WrappedWriter { + return &ipOnlyWriter{dest: dest} +} + +type ipOnlyWriter struct { + dest io.Writer + buf []byte +} + +func (w *ipOnlyWriter) Write(p []byte) (n int, err error) { + if len(p) == 0 { + return 0, nil + } + total := len(p) + w.buf = append(w.buf, p...) + for { + idx := bytes.IndexByte(w.buf, '\n') + if idx == -1 { + break + } + line := w.buf[:idx+1] + // Extract section before first space (header) + headerEnd := bytes.IndexByte(line, ' ') + header := line + if headerEnd >= 0 { + header = line[:headerEnd] + } + // From header, find the first '-' and the next ':' after it + ipStart := bytes.IndexByte(header, '-') + var host []byte + if ipStart >= 0 && ipStart+1 < len(header) { + rest := header[ipStart+1:] + // If IPv6 is bracketed like [2001:db8::1], handle brackets + if len(rest) > 0 && rest[0] == '[' { + // Find closing bracket + rb := bytes.IndexByte(rest, ']') + if rb >= 0 { + host = rest[1:rb] + } + } + if host == nil { + // Take up to the first ':' which precedes the port + colon := bytes.IndexByte(rest, ':') + if colon >= 0 { + host = rest[:colon] + } else { + host = rest + } + } + } + host = bytes.TrimSpace(host) + if host == nil { + host = []byte{} + } + // Write host plus newline + if _, err := w.dest.Write(append(host, '\n')); err != nil { + return 0, err + } + w.buf = w.buf[idx+1:] + } + return total, nil +} + +func (w *ipOnlyWriter) GetOriginal() io.Writer { + return w.dest +} diff --git a/services/util/writer/prefix-writer_test.go b/services/util/writer/prefix-writer_test.go index 4ec278b2..9c5898b1 100644 --- a/services/util/writer/prefix-writer_test.go +++ b/services/util/writer/prefix-writer_test.go @@ -97,3 +97,87 @@ func TestPrefixWriter(t *testing.T) { } }) } + +func TestCleanWriter(t *testing.T) { + for _, tc := range []struct { + desc string + inputs []string + output string + }{ + { + desc: "Basic input with prefix", + inputs: []string{"123-10.0.0.1:9500: some_output_strings\n"}, + output: "some_output_strings\n", + }, + { + desc: "Multiple writes across a single line", + inputs: []string{"123-10.0.0.1:9500:", " some_output_strings\n"}, + output: "some_output_strings\n", + }, + { + desc: "Two lines, both cleaned", + inputs: []string{"1:a b\n2:c d\n"}, + output: "b\nd\n", + }, + { + desc: "Line without space remains unchanged", + inputs: []string{"nospace\n"}, + output: "nospace\n", + }, + } { + t.Run(tc.desc, func(t *testing.T) { + var buf bytes.Buffer + w := GetCleanOutputWriter(&buf) + for _, in := range tc.inputs { + if _, err := w.Write([]byte(in)); err != nil { + t.Fatal(err) + } + } + if got, want := buf.String(), tc.output; got != want { + t.Fatalf("got %q, want %q", got, want) + } + }) + } +} + +func TestIPOnlyWriter(t *testing.T) { + for _, tc := range []struct { + desc string + inputs []string + output string + }{ + { + desc: "IPv4 simple", + inputs: []string{"123-10.0.0.1:9500: foo\n"}, + output: "10.0.0.1\n", + }, + { + desc: "localhost", + inputs: []string{"9-localhost:50042: bar\n"}, + output: "localhost\n", + }, + { + desc: "IPv6 bracketed", + inputs: []string{"77-[2001:db8::1]:50042: baz\n"}, + output: "2001:db8::1\n", + }, + { + desc: "split writes", + inputs: []string{"1-10.1.1.1:50042:", " abc\n"}, + output: "10.1.1.1\n", + }, + } { + t.Run(tc.desc, func(t *testing.T) { + var buf bytes.Buffer + w := GetIPOnlyWriter(&buf) + for _, in := range tc.inputs { + if _, err := w.Write([]byte(in)); err != nil { + t.Fatal(err) + } + } + if got, want := buf.String(), tc.output; got != want { + t.Fatalf("got %q, want %q", got, want) + } + }) + } +}