Add "exec as" feature to public package interface (#516)

services/exec/client/utils.go

@@ -25,14 +25,36 @@ import (
 	pb "github.com/Snowflake-Labs/sansshell/services/exec"
 )
 
+// ExecRemoteCommand is an options to exec a command on a remote host
+type execRemoteOptions 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 ExecRemoteOption func(*execRemoteOptions)
+
+func ExecAsUser(user string) ExecRemoteOption {
+	return func(o *execRemoteOptions) {
+		o.execAsUser = &user
+	}
+}
+
 // ExecRemoteCommand 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 ExecRemoteCommandWithOpts instead.
 func ExecRemoteCommand(ctx context.Context, conn *proxy.Conn, binary string, args ...string) (*pb.ExecResponse, error) {
+	return ExecRemoteCommandWithOpts(ctx, conn, binary, args)
+}
+
+// ExecRemoteCommandWithOpts 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 ExecRemoteCommandWithOpts(ctx context.Context, conn *proxy.Conn, binary string, args []string, opts ...ExecRemoteOption) (*pb.ExecResponse, error) {
 	if len(conn.Targets) != 1 {
 		return nil, errors.New("ExecRemoteCommand only supports single targets")
 	}
 
-	result, err := ExecRemoteCommandMany(ctx, conn, binary, args...)
+	result, err := ExecRemoteCommandManyWithOpts(ctx, conn, binary, args)
 	if err != nil {
 		return nil, err
 	}
@@ -45,15 +67,34 @@ func ExecRemoteCommand(ctx context.Context, conn *proxy.Conn, binary string, arg
 	return result[0].Resp, nil
 }
 
-// ExecRemoteCommand is a helper function for execing a command on one or remote hosts
+// ExecRemoteCommandMany 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 ExecRemoteCommandManyWithOpts instead.
 func ExecRemoteCommandMany(ctx context.Context, conn *proxy.Conn, binary string, args ...string) ([]*pb.RunManyResponse, error) {
+	return ExecRemoteCommandManyWithOpts(ctx, conn, binary, args)
+}
+
+// ExecRemoteCommandManyWithOpts 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 ExecRemoteCommandManyWithOpts(ctx context.Context, conn *proxy.Conn, binary string, args []string, opts ...ExecRemoteOption) ([]*pb.RunManyResponse, error) {
+	execOptions := &execRemoteOptions{}
+	for _, opt := range opts {
+		opt(execOptions)
+	}
+
 	c := pb.NewExecClientProxy(conn)
 	req := &pb.ExecRequest{
 		Command: binary,
 		Args:    args,
 	}
+
+	if execOptions.execAsUser != nil {
+		req.User = *execOptions.execAsUser
+	}
+
 	respChan, err := c.RunOneMany(ctx, req)
 	if err != nil {
 		return nil, err

services/sansshell/client/utils.go

@@ -0,0 +1,47 @@
+/* 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"
+
+	"google.golang.org/protobuf/types/known/emptypb"
+
+	"github.com/Snowflake-Labs/sansshell/proxy/proxy"
+	pb "github.com/Snowflake-Labs/sansshell/services/sansshell"
+)
+
+// GetSansshellVersion gets the version of the sansshell server for a particular target.
+// Returns error if:
+// - was provided not exactly one target
+// - could not get the version from target
+func GetSansshellVersion(ctx context.Context, conn *proxy.Conn) (string, error) {
+	if len(conn.Targets) != 1 {
+		return "", errors.New("GetSansshellVersion support only single target")
+	}
+
+	client := pb.NewStateClient(conn)
+
+	resp, err := client.Version(ctx, &emptypb.Empty{})
+	if err != nil {
+		return "", fmt.Errorf("Could not get proxy version: %v\n", err)
+	}
+
+	return resp.GetVersion(), nil
+}