gnugmanov SNOW 2334154 sanssh clean output

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 '<index>-<target>: '
 	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
@@ -313,6 +317,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(

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 '<index>-<target>: '")
-	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 '<index>-<target>: '")
+	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,
 	}

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
+}

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
+// "<random>-<host-or-ip>:<port>: ..." 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
+}

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)
+			}
+		})
+	}
+}