SNOW-1991908: add fdbexec

cmd/sanssh/main.go

@@ -45,6 +45,7 @@ import (
 	_ "github.com/Snowflake-Labs/sansshell/services/dns/client"
 	_ "github.com/Snowflake-Labs/sansshell/services/exec/client"
 	_ "github.com/Snowflake-Labs/sansshell/services/fdb/client"
+	_ "github.com/Snowflake-Labs/sansshell/services/fdbexec/client"
 	_ "github.com/Snowflake-Labs/sansshell/services/healthcheck/client"
 	_ "github.com/Snowflake-Labs/sansshell/services/httpoverrpc/client"
 	_ "github.com/Snowflake-Labs/sansshell/services/localfile/client"

cmd/sansshell-server/default-policy.rego

@@ -33,6 +33,9 @@ allow {
 	input.message.filename = "/etc/hosts"
 }
 
+allow {
+	input.type = "FdbExec.FdbExecRequest"
+}
 allow {
 	input.type = "Exec.ExecRequest"
 	input.message.command = "/bin/echo"

cmd/sansshell-server/main.go

@@ -59,6 +59,7 @@ import (
 	ansible "github.com/Snowflake-Labs/sansshell/services/ansible/server"
 	_ "github.com/Snowflake-Labs/sansshell/services/dns/server"
 	_ "github.com/Snowflake-Labs/sansshell/services/exec/server"
+	_ "github.com/Snowflake-Labs/sansshell/services/fdbexec/server"
 	_ "github.com/Snowflake-Labs/sansshell/services/httpoverrpc/server"
 	_ "github.com/Snowflake-Labs/sansshell/services/tlsinfo/server"
 

services/fdbexec/README.md

@@ -0,0 +1,17 @@
+# FdbExec
+
+The FdbExec service provides a way to execute commands on remote hosts.
+
+## Usage
+
+### sanssh fdbexec run
+```
+sanssh <sanssh-args> fdbexec run [--stream] [--user user] <command> [<args>...]
+```
+
+## Notes
+
+The FdbExec service is identical to the Exec service in functionality, but is specifically intended for use by the FDB team. The rego policy can be configured to allow only FDB team members to use this command.
+
+The `--stream` flag streams command output as it is generated.
+The `--user user` flag executes the command as the specified user (requires privileges). 

services/fdbexec/client/client.go

@@ -0,0 +1,169 @@
+/* Copyright (c) 2019 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 client provides the client interface for 'fdbexec'
+package client
+
+import (
+	"context"
+	"flag"
+	"fmt"
+	"io"
+	"os"
+
+	"github.com/google/subcommands"
+
+	"github.com/Snowflake-Labs/sansshell/client"
+	pb "github.com/Snowflake-Labs/sansshell/services/fdbexec"
+	"github.com/Snowflake-Labs/sansshell/services/util"
+)
+
+// Default value for using streaming exec.
+// This is exposed so that clients can set a default that works for
+// their environment. StreamingExec was added after Exec, so policies
+// or sansshell nodes may not have been updated to support streaming.
+var DefaultStreaming = false
+
+const subPackage = "fdbexec"
+
+func init() {
+	subcommands.Register(&fdbExecCmd{}, subPackage)
+}
+
+func (*fdbExecCmd) GetSubpackage(f *flag.FlagSet) *subcommands.Commander {
+	c := client.SetupSubpackage(subPackage, f)
+	c.Register(&runCmd{}, "")
+	return c
+}
+
+type fdbExecCmd struct{}
+
+func (*fdbExecCmd) Name() string { return subPackage }
+func (p *fdbExecCmd) Synopsis() string {
+	return client.GenerateSynopsis(p.GetSubpackage(flag.NewFlagSet("", flag.ContinueOnError)), 2)
+}
+func (p *fdbExecCmd) Usage() string {
+	return client.GenerateUsage(subPackage, p.Synopsis())
+}
+func (*fdbExecCmd) SetFlags(f *flag.FlagSet) {}
+
+func (p *fdbExecCmd) Execute(ctx context.Context, f *flag.FlagSet, args ...interface{}) subcommands.ExitStatus {
+	c := p.GetSubpackage(f)
+	return c.Execute(ctx, args...)
+}
+
+type runCmd struct {
+	streaming bool
+	user      string
+
+	// returnCode internally keeps track of the final status to return
+	returnCode subcommands.ExitStatus
+}
+
+func (*runCmd) Name() string     { return "run" }
+func (*runCmd) Synopsis() string { return "Run provided command and return a response." }
+func (*runCmd) Usage() string {
+	return `run [--stream] [--user=user] <command> [<args>...]:
+  Run a command remotely and return the response
+
+	Note: This is not optimized for large output or long running commands.  If
+	the output doesn't fit in memory in a single proto message or if it doesn't
+	complete within the timeout, you'll have a bad time.
+
+	The --stream flag can be used to stream back command output as the command
+	runs. It doesn't affect the timeout.
+
+	--user flag allows to specify a user for running command, equivalent of
+	sudo -u <user> <command> ...
+`
+}
+
+func (p *runCmd) SetFlags(f *flag.FlagSet) {
+	f.BoolVar(&p.streaming, "stream", DefaultStreaming, "If true, stream back stdout and stdin during the command instead of sending it all at the end.")
+	f.StringVar(&p.user, "user", "", "If specified, allows to run a command as a specified user. Equivalent of sudo -u <user> <command> ... .")
+}
+
+func (p *runCmd) printCommandOutput(state *util.ExecuteState, idx int, resp *pb.FdbExecResponse, err error) {
+	if err == io.EOF {
+		// Streaming commands may return EOF
+		return
+	}
+	if err != nil {
+		fmt.Fprintf(state.Err[idx], "Command execution failure - %v\n", err)
+		// If any target had errors it needs to be reported for that target but we still
+		// need to process responses off the channel. Final return code though should
+		// indicate something failed.
+		p.returnCode = subcommands.ExitFailure
+		return
+	}
+	if len(resp.Stderr) > 0 {
+		fmt.Fprintf(state.Err[idx], "%s", resp.Stderr)
+	}
+	fmt.Fprintf(state.Out[idx], "%s", resp.Stdout)
+	if resp.RetCode != 0 {
+		p.returnCode = subcommands.ExitFailure
+	}
+}
+
+func (p *runCmd) Execute(ctx context.Context, f *flag.FlagSet, args ...interface{}) subcommands.ExitStatus {
+	state := args[0].(*util.ExecuteState)
+	if f.NArg() == 0 {
+		fmt.Fprintf(os.Stderr, "Please specify a command to execute.\n")
+		return subcommands.ExitUsageError
+	}
+
+	c := pb.NewFdbExecClientProxy(state.Conn)
+	req := &pb.FdbExecRequest{Command: f.Args()[0], Args: f.Args()[1:], User: p.user}
+
+	if p.streaming {
+		resp, err := c.StreamingRunOneMany(ctx, req)
+		if err != nil {
+			// Emit this to every error file as it's not specific to a given target.
+			for _, e := range state.Err {
+				fmt.Fprintf(e, "All targets - could not execute: %v\n", err)
+			}
+			return subcommands.ExitFailure
+		}
+
+		for {
+			rs, err := resp.Recv()
+			if err != nil {
+				if err == io.EOF {
+					return p.returnCode
+				}
+				fmt.Fprintf(os.Stderr, "Stream failure: %v\n", err)
+				return subcommands.ExitFailure
+			}
+			for _, r := range rs {
+				p.printCommandOutput(state, r.Index, r.Resp, r.Error)
+			}
+		}
+	} else {
+		resp, err := c.RunOneMany(ctx, req)
+		if err != nil {
+			// Emit this to every error file as it's not specific to a given target.
+			for _, e := range state.Err {
+				fmt.Fprintf(e, "All targets - could not execute: %v\n", err)
+			}
+			return subcommands.ExitFailure
+		}
+
+		for r := range resp {
+			p.printCommandOutput(state, r.Index, r.Resp, r.Error)
+		}
+	}
+	return p.returnCode
+}

services/fdbexec/client/utils.go

@@ -0,0 +1,108 @@
+/* Copyright (c) 2019 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 client
+
+import (
+	"context"
+	"errors"
+	"fmt"
+
+	"github.com/Snowflake-Labs/sansshell/proxy/proxy"
+	pb "github.com/Snowflake-Labs/sansshell/services/fdbexec"
+)
+
+// FdbExecRemoteCommand is an options to exec a command on a remote host
+type fdbExecRemoteOptions struct {
+	// execAsUser, Optional. User is the user to run the command as, if nil the command will be run as the default user
+	execAsUser *string
+}
+
+type FdbExecRemoteOption func(*fdbExecRemoteOptions)
+
+func FdbExecAsUser(user string) FdbExecRemoteOption {
+	return func(o *fdbExecRemoteOptions) {
+		o.execAsUser = &user
+	}
+}
+
+// FdbExecRemoteCommand is a helper function for execing a command on a remote host
+// using a proxy.Conn. If the conn is defined for >1 targets this will return an error.
+//
+// Deprecated: Use FdbExecRemoteCommandWithOpts instead.
+func FdbExecRemoteCommand(ctx context.Context, conn *proxy.Conn, binary string, args ...string) (*pb.FdbExecResponse, error) {
+	return FdbExecRemoteCommandWithOpts(ctx, conn, binary, args)
+}
+
+// FdbExecRemoteCommandWithOpts is a helper function for execing a command on a remote host with provided opts
+// using a proxy.Conn. If the conn is defined for >1 targets this will return an error.
+func FdbExecRemoteCommandWithOpts(ctx context.Context, conn *proxy.Conn, binary string, args []string, opts ...FdbExecRemoteOption) (*pb.FdbExecResponse, error) {
+	if len(conn.Targets) != 1 {
+		return nil, errors.New("FdbExecRemoteCommand only supports single targets")
+	}
+
+	result, err := FdbExecRemoteCommandManyWithOpts(ctx, conn, binary, args)
+	if err != nil {
+		return nil, err
+	}
+	if len(result) < 1 {
+		return nil, fmt.Errorf("FdbExecRemoteCommand error: received empty response")
+	}
+	if result[0].Error != nil {
+		return nil, fmt.Errorf("FdbExecRemoteCommand error: %v", result[0].Error)
+	}
+	return result[0].Resp, nil
+}
+
+// FdbExecRemoteCommandMany is a helper function for execing a command on one or remote hosts
+// using a proxy.Conn.
+// `binary` refers to the absolute path of the binary file on the remote host.
+//
+// // Deprecated: Use FdbExecRemoteCommandManyWithOpts instead.
+func FdbExecRemoteCommandMany(ctx context.Context, conn *proxy.Conn, binary string, args ...string) ([]*pb.RunManyResponse, error) {
+	return FdbExecRemoteCommandManyWithOpts(ctx, conn, binary, args)
+}
+
+// FdbExecRemoteCommandManyWithOpts is a helper function for execing a command on one or remote hosts with provided opts
+// using a proxy.Conn.
+// `binary` refers to the absolute path of the binary file on the remote host.
+func FdbExecRemoteCommandManyWithOpts(ctx context.Context, conn *proxy.Conn, binary string, args []string, opts ...FdbExecRemoteOption) ([]*pb.RunManyResponse, error) {
+	execOptions := &fdbExecRemoteOptions{}
+	for _, opt := range opts {
+		opt(execOptions)
+	}
+
+	c := pb.NewFdbExecClientProxy(conn)
+	req := &pb.FdbExecRequest{
+		Command: binary,
+		Args:    args,
+	}
+
+	if execOptions.execAsUser != nil {
+		req.User = *execOptions.execAsUser
+	}
+
+	respChan, err := c.RunOneMany(ctx, req)
+	if err != nil {
+		return nil, err
+	}
+	result := make([]*pb.RunManyResponse, len(conn.Targets))
+	for r := range respChan {
+		result[r.Index] = r
+	}
+
+	return result, nil
+}

services/fdbexec/fdbexec.go

@@ -0,0 +1,22 @@
+/* Copyright (c) 2019 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 fdbexec defines the RPC interface for the sansshell FdbExec actions.
+package fdbexec
+
+// To regenerate the proto headers if the .proto changes, just run go generate
+// This comment encodes the necessary magic:
+//go:generate protoc --go_out=. --go_opt=paths=source_relative --go-grpc_out=require_unimplemented_servers=false:. --go-grpc_opt=paths=source_relative --go-grpcproxy_out=. --go-grpcproxy_opt=paths=source_relative fdbexec.proto