From dc963f5e480bd3a64077792ce1fa41f3e54a2d0f Mon Sep 17 00:00:00 2001 From: Ivan Velichko Date: Fri, 27 Sep 2024 13:10:44 +0000 Subject: [PATCH] add labctl cp command --- cmd/cp/cp.go | 149 +++++++++++++++++++++++++++++++++++++++ cmd/sshproxy/sshproxy.go | 86 +++++++++++++++------- go.sum | 6 -- main.go | 2 + 4 files changed, 212 insertions(+), 31 deletions(-) create mode 100644 cmd/cp/cp.go diff --git a/cmd/cp/cp.go b/cmd/cp/cp.go new file mode 100644 index 0000000..1ac4c7e --- /dev/null +++ b/cmd/cp/cp.go @@ -0,0 +1,149 @@ +package cp + +import ( + "context" + "fmt" + "os/exec" + "strings" + + "github.com/spf13/cobra" + + "github.com/iximiuz/labctl/cmd/sshproxy" + "github.com/iximiuz/labctl/internal/labcli" +) + +const example = ` # Copy a file to the playground + labctl cp 65e78a64366c2b0cf9ddc34c:/home/laborant/some/file ./some/file + + # Copy a file from the playground + labctl cp ./some/file 65e78a64366c2b0cf9ddc34c:/home/laborant/some/file + + # Copy a directory to the playground + labctl cp -r ./some/dir 65e78a64366c2b0cf9ddc34c:/home/laborant/some/dir + + # Copy a directory from the playground + labctl cp 65e78a64366c2b0cf9ddc34c:/home/laborant/some/dir ./some/dir +` + +type Direction string + +const ( + DirectionLocalToRemote Direction = "local-to-remote" + DirectionRemoteToLocal Direction = "remote-to-local" +) + +type options struct { + machine string + user string + + playID string + localPath string + remotePath string + recursive bool + + direction Direction +} + +func NewCommand(cli labcli.CLI) *cobra.Command { + var opts options + + cmd := &cobra.Command{ + Use: "cp [flags] : \n labctl cp [flags] :", + Short: `Copy files to and from the target playground`, + Example: example, + Args: cobra.MinimumNArgs(2), + RunE: func(cmd *cobra.Command, args []string) error { + if strings.Contains(args[0], ":") { + parts := strings.Split(args[0], ":") + opts.playID = parts[0] + opts.remotePath = parts[1] + + opts.localPath = args[1] + opts.direction = DirectionRemoteToLocal + } else { + parts := strings.Split(args[1], ":") + opts.playID = parts[0] + opts.remotePath = parts[1] + + opts.localPath = args[0] + opts.direction = DirectionLocalToRemote + } + + return labcli.WrapStatusError(runCopy(cmd.Context(), cli, &opts)) + }, + } + + flags := cmd.Flags() + + flags.StringVarP( + &opts.machine, + "machine", + "m", + "", + `Target machine (default: the first machine in the playground)`, + ) + flags.StringVarP( + &opts.user, + "user", + "u", + "", + `SSH user (default: the machine's default login user)`, + ) + flags.BoolVarP( + &opts.recursive, + "recursive", + "r", + false, + `Copy directories recursively`, + ) + + return cmd +} + +func runCopy(ctx context.Context, cli labcli.CLI, opts *options) error { + ctx, cancel := context.WithCancel(ctx) + defer cancel() + + return sshproxy.RunSSHProxy(ctx, cli, &sshproxy.Options{ + PlayID: opts.playID, + Machine: opts.machine, + User: opts.user, + Quiet: true, + WithProxy: func(ctx context.Context, info *sshproxy.SSHProxyInfo) error { + args := []string{ + "-i", info.IdentityFile, + "-o", "StrictHostKeyChecking=no", + "-o", "UserKnownHostsFile=/dev/null", + "-P", info.ProxyPort, + "-C", // compress + } + + if opts.recursive { + args = append(args, "-r") + } + + if opts.direction == DirectionLocalToRemote { + args = append(args, + opts.localPath, + fmt.Sprintf("%s@%s:%s", info.User, info.ProxyHost, opts.remotePath), + ) + } else { + args = append(args, + fmt.Sprintf("%s@%s:%s", info.User, info.ProxyHost, opts.remotePath), + opts.localPath, + ) + } + + cmd := exec.CommandContext(ctx, "scp", args...) + cmd.Stdout = cli.OutputStream() + cmd.Stderr = cli.ErrorStream() + + if err := cmd.Run(); err != nil { + return fmt.Errorf("copy command failed %s: %w", cmd.String(), err) + } + + cli.PrintAux("Done!\n") + return nil + }, + }) +} diff --git a/cmd/sshproxy/sshproxy.go b/cmd/sshproxy/sshproxy.go index 73bc435..c13f58a 100644 --- a/cmd/sshproxy/sshproxy.go +++ b/cmd/sshproxy/sshproxy.go @@ -5,6 +5,7 @@ import ( "fmt" "log/slog" "os/exec" + "path/filepath" "strings" "github.com/spf13/cobra" @@ -20,7 +21,10 @@ type Options struct { User string Address string - IDE bool + IDE bool + Quiet bool + + WithProxy func(ctx context.Context, info *SSHProxyInfo) error } func NewCommand(cli labcli.CLI) *cobra.Command { @@ -31,6 +35,8 @@ func NewCommand(cli labcli.CLI) *cobra.Command { Short: `Start SSH proxy to the playground's machine`, Args: cobra.ExactArgs(1), RunE: func(cmd *cobra.Command, args []string) error { + cli.SetQuiet(opts.Quiet) + opts.PlayID = args[0] if opts.Address != "" && strings.Count(opts.Address, ":") != 1 { @@ -68,10 +74,25 @@ func NewCommand(cli labcli.CLI) *cobra.Command { false, `Open the playground in the IDE (only VSCode is supported at the moment)`, ) + flags.BoolVarP( + &opts.Quiet, + "quiet", + "q", + false, + `Quiet mode (don't print any messages except errors)`, + ) return cmd } +type SSHProxyInfo struct { + User string + Machine string + ProxyHost string + ProxyPort string + IdentityFile string +} + func RunSSHProxy(ctx context.Context, cli labcli.CLI, opts *Options) error { p, err := cli.Client().GetPlay(ctx, opts.PlayID) if err != nil { @@ -140,28 +161,7 @@ func RunSSHProxy(ctx context.Context, cli labcli.CLI, opts *Options) error { } }() - if !opts.IDE { - cli.PrintOut("SSH proxy is running on %s\n", localPort) - cli.PrintOut( - "\n# Connect from the terminal:\nssh -i %s/%s ssh://%s@%s:%s\n", - cli.Config().SSHDir, ssh.IdentityFile, opts.User, localHost, localPort, - ) - - cli.PrintOut("\n# Or add the following to your ~/.ssh/config:\n") - cli.PrintOut("Host %s\n", opts.PlayID+"-"+opts.Machine) - cli.PrintOut(" HostName %s\n", localHost) - cli.PrintOut(" Port %s\n", localPort) - cli.PrintOut(" User %s\n", opts.User) - cli.PrintOut(" IdentityFile %s/%s\n", cli.Config().SSHDir, ssh.IdentityFile) - cli.PrintOut(" StrictHostKeyChecking no\n") - cli.PrintOut(" UserKnownHostsFile /dev/null\n\n") - - cli.PrintOut("# To access the playground in Visual Studio Code:\n") - cli.PrintOut("code --folder-uri vscode-remote://ssh-remote+%s@%s:%s%s\n\n", - opts.User, localHost, localPort, userHomeDir(opts.User)) - - cli.PrintOut("\nPress Ctrl+C to stop\n") - } else { + if opts.IDE { cli.PrintAux("Opening the playground in the IDE...\n") // Hack: SSH into the playground first - otherwise, VSCode will fail to connect for some reason. @@ -184,8 +184,44 @@ func RunSSHProxy(ctx context.Context, cli labcli.CLI, opts *Options) error { } } - // Wait for ctrl+c - <-ctx.Done() + if !opts.IDE && !opts.Quiet { + cli.PrintAux("SSH proxy is running on %s\n", localPort) + cli.PrintAux( + "\n# Connect from the terminal:\nssh -i %s/%s ssh://%s@%s:%s\n", + cli.Config().SSHDir, ssh.IdentityFile, opts.User, localHost, localPort, + ) + + cli.PrintAux("\n# Or add the following to your ~/.ssh/config:\n") + cli.PrintAux("Host %s\n", opts.PlayID+"-"+opts.Machine) + cli.PrintAux(" HostName %s\n", localHost) + cli.PrintAux(" Port %s\n", localPort) + cli.PrintAux(" User %s\n", opts.User) + cli.PrintAux(" IdentityFile %s/%s\n", cli.Config().SSHDir, ssh.IdentityFile) + cli.PrintAux(" StrictHostKeyChecking no\n") + cli.PrintAux(" UserKnownHostsFile /dev/null\n\n") + + cli.PrintAux("# To access the playground in Visual Studio Code:\n") + cli.PrintAux("code --folder-uri vscode-remote://ssh-remote+%s@%s:%s%s\n\n", + opts.User, localHost, localPort, userHomeDir(opts.User)) + + cli.PrintAux("\nPress Ctrl+C to stop\n") + } + + if opts.WithProxy != nil { + info := &SSHProxyInfo{ + User: opts.User, + Machine: opts.Machine, + ProxyHost: localHost, + ProxyPort: localPort, + IdentityFile: filepath.Join(cli.Config().SSHDir, ssh.IdentityFile), + } + if err := opts.WithProxy(ctx, info); err != nil { + return fmt.Errorf("proxy callback failed: %w", err) + } + } else { + // Wait for ctrl+c + <-ctx.Done() + } return nil } diff --git a/go.sum b/go.sum index 414e958..c858e9a 100644 --- a/go.sum +++ b/go.sum @@ -18,12 +18,8 @@ github.com/charmbracelet/huh v0.6.0 h1:mZM8VvZGuE0hoDXq6XLxRtgfWyTI3b2jZNKh0xWma github.com/charmbracelet/huh v0.6.0/go.mod h1:GGNKeWCeNzKpEOh/OJD8WBwTQjV3prFAtQPpLv+AVwU= github.com/charmbracelet/lipgloss v0.13.0 h1:4X3PPeoWEDCMvzDvGmTajSyYPcZM4+y8sCA/SsA3cjw= github.com/charmbracelet/lipgloss v0.13.0/go.mod h1:nw4zy0SBX/F/eAO1cWdcvy6qnkDUxr8Lw7dvFrAIbbY= -github.com/charmbracelet/x/ansi v0.3.1 h1:CRO6lc/6HCx2/D6S/GZ87jDvRvk6GtPyFP+IljkNtqI= -github.com/charmbracelet/x/ansi v0.3.1/go.mod h1:dk73KoMTT5AX5BsX0KrqhsTqAnhZZoCBjs7dGWp4Ktw= github.com/charmbracelet/x/ansi v0.3.2 h1:wsEwgAN+C9U06l9dCVMX0/L3x7ptvY1qmjMwyfE6USY= github.com/charmbracelet/x/ansi v0.3.2/go.mod h1:dk73KoMTT5AX5BsX0KrqhsTqAnhZZoCBjs7dGWp4Ktw= -github.com/charmbracelet/x/exp/strings v0.0.0-20240914193755-48d9a4a13687 h1:tAmXpXce4cvIGUE0A6qKs/Q+vyUREBqdrlLjGjTHMr4= -github.com/charmbracelet/x/exp/strings v0.0.0-20240914193755-48d9a4a13687/go.mod h1:pBhA0ybfXv6hDjQUZ7hk1lVxBiUbupdw5R31yPUViVQ= github.com/charmbracelet/x/exp/strings v0.0.0-20240919170804-a4978c8e603a h1:JMdM89Udp/cOl5tC3MuUJXTPE/nAdU1oyt9jRU44qq8= github.com/charmbracelet/x/exp/strings v0.0.0-20240919170804-a4978c8e603a/go.mod h1:pBhA0ybfXv6hDjQUZ7hk1lVxBiUbupdw5R31yPUViVQ= github.com/charmbracelet/x/term v0.2.0 h1:cNB9Ot9q8I711MyZ7myUR5HFWL/lc3OpU8jZ4hwm0x0= @@ -34,8 +30,6 @@ github.com/creack/pty v1.1.18/go.mod h1:MOBLtS5ELjhRRrroQr9kyvTxUAFNvYEK993ew/Vr github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= -github.com/docker/cli v27.2.1+incompatible h1:U5BPtiD0viUzjGAjV1p0MGB8eVA3L3cbIrnyWmSJI70= -github.com/docker/cli v27.2.1+incompatible/go.mod h1:JLrzqnKDaYBop7H2jaqPtU4hHvMKP+vjCwu2uszcLI8= github.com/docker/cli v27.3.1+incompatible h1:qEGdFBF3Xu6SCvCYhc7CzaQTlBmqDuzxPDpigSyeKQQ= github.com/docker/cli v27.3.1+incompatible/go.mod h1:JLrzqnKDaYBop7H2jaqPtU4hHvMKP+vjCwu2uszcLI8= github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= diff --git a/main.go b/main.go index 3cd8651..dd0a26a 100644 --- a/main.go +++ b/main.go @@ -11,6 +11,7 @@ import ( "github.com/iximiuz/labctl/cmd/auth" "github.com/iximiuz/labctl/cmd/challenge" "github.com/iximiuz/labctl/cmd/content" + "github.com/iximiuz/labctl/cmd/cp" "github.com/iximiuz/labctl/cmd/playground" "github.com/iximiuz/labctl/cmd/portforward" "github.com/iximiuz/labctl/cmd/ssh" @@ -70,6 +71,7 @@ func main() { auth.NewCommand(cli), challenge.NewCommand(cli), content.NewCommand(cli), + cp.NewCommand(cli), playground.NewCommand(cli), portforward.NewCommand(cli), ssh.NewCommand(cli),