Skip to content

Commit

Permalink
add labctl cp command
Browse files Browse the repository at this point in the history
  • Loading branch information
iximiuz committed Sep 27, 2024
1 parent 69acfff commit dc963f5
Show file tree
Hide file tree
Showing 4 changed files with 212 additions and 31 deletions.
149 changes: 149 additions & 0 deletions cmd/cp/cp.go
Original file line number Diff line number Diff line change
@@ -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] <playground-id>:<source-path> <destination-path>\n labctl cp [flags] <source-path> <playground-id>:<destination-path>",
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
},
})
}
86 changes: 61 additions & 25 deletions cmd/sshproxy/sshproxy.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import (
"fmt"
"log/slog"
"os/exec"
"path/filepath"
"strings"

"github.com/spf13/cobra"
Expand All @@ -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 {
Expand All @@ -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 {
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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.
Expand All @@ -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
}
Expand Down
6 changes: 0 additions & 6 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand All @@ -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=
Expand Down
2 changes: 2 additions & 0 deletions main.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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),
Expand Down

0 comments on commit dc963f5

Please sign in to comment.