Merge pull request #18 from adzimzf/improve_forwarding

Improve forwarding

main.go

@@ -6,6 +6,7 @@ import (
 	"io"
 	"log"
 	"net"
+	"strings"
 	"time"
 
 	"github.com/adzimzf/tpot/config"
@@ -40,14 +41,15 @@ func main() {
 }
 
 const example = `
-tpot -c --add         // Set up the configuration environment
-tpot -c --edit        // Edit all the configuration
-tpot staging          // Show the node list of staging environment
-tpot staging --edit   // Edit the staging proxy configuration
-tpot prod -a          // Get the latest node list then append to the cache for production 
-tpot prod -r          // Refresh the cache with the latest node from Teleport UI
-tpot prod -u root     // Login into production using root user
-tpot prod -L          // Run the tsh forwarding based on the config list
+tpot -c --add                       // Set up the configuration environment
+tpot -c --edit                      // Edit all the configuration
+tpot staging                        // Show the node list of staging environment
+tpot staging --edit                 // Edit the staging proxy configuration
+tpot prod -a                        // Get the latest node list then append to the cache for production 
+tpot prod -r                        // Refresh the cache with the latest node from Teleport UI
+tpot prod -u root                   // Login into production using root user
+tpot prod -L                        // Run the tsh forwarding based on the config list
+tpot prod -L 123:localhost:123      // Run the tsh forwarding based on the list in argument
 `
 
 var rootCmd = &cobra.Command{
@@ -96,6 +98,28 @@ var rootCmd = &cobra.Command{
 			}
 			proxy.Node = *node
 
+			forwardingNodes := proxy.Forwarding.Nodes
+			if len(args) > 1 {
+				nodesStr := strings.Split(args[1], ",")
+				nodes := []*config.ForwardingNode{}
+				for _, s := range nodesStr {
+					parts := strings.Split(s, ":")
+					if len(parts) != 3 {
+						cmd.PrintErrln(fmt.Sprintf("invalid forwarding format for: %s, use format <local port>:<remote address>:<remote port> example: 123:localhost:123", s))
+						return
+					}
+					nodes = append(nodes, &config.ForwardingNode{
+						ListenPort: parts[0],
+						RemoteHost: parts[1],
+						RemotePort: parts[2],
+					})
+				}
+				// replace the forwarding config nodes
+				if len(nodes) > 0 {
+					forwardingNodes = nodes
+				}
+			}
+
 			host := ui.GetSelectedHost(proxy.Node.ListHostname())
 			if host == "" {
 				cmd.PrintErrln("Pick at least one host to login")
@@ -110,7 +134,7 @@ var rootCmd = &cobra.Command{
 
 			f := fwd{
 				tsh:         tsh.NewTSH(proxy),
-				list:        proxy.Forwarding.Nodes,
+				list:        forwardingNodes,
 				nodeHost:    host,
 				defaultUser: user,
 			}
@@ -402,6 +426,12 @@ func (f *fwd) Run() error {
 	if len(f.list) == 0 {
 		return fmt.Errorf("forwarding configuration is empty")
 	}
+
+	err := f.tsh.Login()
+	if err != nil {
+		return fmt.Errorf("failed to login, error: %v", err)
+	}
+
 	for _, node := range f.list {
 		go func(node *config.ForwardingNode) {
 			f.execForwarding(node)
@@ -414,6 +444,7 @@ func (f *fwd) Run() error {
 
 func (f *fwd) doHealthCheck() {
 	for {
+		time.Sleep(2 * time.Second)
 		for _, node := range f.list {
 			go func(node *config.ForwardingNode) {
 				timeout := time.Second
@@ -428,7 +459,6 @@ func (f *fwd) doHealthCheck() {
 				node.Error = ""
 			}(node)
 		}
-		time.Sleep(2 * time.Second)
 	}
 }
 
@@ -463,6 +493,8 @@ type sleepReader struct {
 
 // Read will only delay the reader then return error
 func (f *sleepReader) Read(p []byte) (n int, err error) {
-	time.Sleep(f.dur)
+	for {
+		time.Sleep(2 * time.Second)
+	}
 	return 0, io.EOF
 }

tsh/runner.go

@@ -83,7 +83,7 @@ func (t *TSH) SSH(username, host string) error {
 // ListNodes get the list nodes from proxy
 func (t *TSH) ListNodes() (config.Node, error) {
 
-	if err := t.login(); err != nil {
+	if err := t.Login(); err != nil {
 		return config.Node{}, err
 	}
 
@@ -195,7 +195,7 @@ func trimSliceString(list []string) (res []string) {
 	return
 }
 
-func (t *TSH) login() error {
+func (t *TSH) Login() error {
 
 	if t.isLogin() {
 		return nil

ui/hostname_table.go

@@ -178,7 +178,7 @@ func findMaxXY(s string) (m maxXY) {
 
 }
 
-func debug(i ...interface{}) {
+func Debug(i ...interface{}) {
 	s := time.Now().String() + "\n"
 	for _, i1 := range i {
 		s += fmt.Sprintf("%v", i1)