Merge pull request #4 from TNK-Studio/dev

[ADD] add edit server menu

Author: elfgzp

README.md

@@ -1,2 +1,3 @@
 # gortal
-A SSH Portal written by Go, use to manage & jump to any servers.
+
+A lightweight jumpserver written by Go, use to manage & jump to any servers.

config/config.go

@@ -6,6 +6,7 @@ import (
 	"log"
 	"os"
 
+	"github.com/TNK-Studio/gortal/utils/logger"
 	"gopkg.in/yaml.v2"
 )
 
@@ -56,24 +57,24 @@ type SSHUser struct {
 func (c *Config) ReadFrom(path string) error {
 	configFile, err := ioutil.ReadFile(path)
 	if err != nil {
-		fmt.Printf("Error reading YAML file: %s\n", err)
+		logger.Logger.Infof("Error reading YAML file: %s\n", err)
 		return err
 	}
 
 	err = yaml.Unmarshal([]byte(configFile), c)
 	if err != nil {
-		fmt.Printf("Error parsing YAML file: %s\n", err)
+		logger.Logger.Infof("Error parsing YAML file: %s\n", err)
 		return err
 	}
 	return nil
 }
 
 // SaveTo save config
 func (c *Config) SaveTo(path string) error {
-	fmt.Printf("Save config to '%s'\n", path)
+	logger.Logger.Infof("Save config to '%s'\n", path)
 	bytes, err := yaml.Marshal(c)
 	if err != nil {
-		fmt.Printf("Error parsing YAML obj: %s\n", err)
+		logger.Logger.Infof("Error parsing YAML obj: %s\n", err)
 		return err
 	}
 	ioutil.WriteFile(path, bytes, 0644)
@@ -151,9 +152,9 @@ func (c *Config) GetUserByUsername(username string) *User {
 // }
 
 // GetUserServers get user servers list
-func (c *Config) GetUserServers(user *User) []*Server {
-	servers := make([]*Server, 0)
-	for _, server := range *c.Servers {
+func (c *Config) GetUserServers(user *User) map[string]*Server {
+	servers := make(map[string]*Server, 0)
+	for serverKey, server := range *c.Servers {
 	loop:
 		for _, sshUser := range *server.SSHUsers {
 			if sshUser.AllowUsers == nil {
@@ -166,26 +167,23 @@ func (c *Config) GetUserServers(user *User) []*Server {
 				}
 			}
 		}
-		servers = append(
-			servers,
-			server,
-		)
+		servers[serverKey] = server
 	}
 	return servers
 }
 
 // GetServerSSHUsers get all allow server' s ssh users
-func (c *Config) GetServerSSHUsers(user *User, server *Server) []*SSHUser {
-	sshUsers := make([]*SSHUser, 0)
-	for _, sshUser := range *server.SSHUsers {
+func (c *Config) GetServerSSHUsers(user *User, server *Server) map[string]*SSHUser {
+	sshUsers := make(map[string]*SSHUser, 0)
+	for sshUserKey, sshUser := range *server.SSHUsers {
 		if sshUser.AllowUsers == nil {
-			sshUsers = append(sshUsers, sshUser)
+			sshUsers[sshUserKey] = sshUser
 			continue
 		}
 
 		for _, username := range *sshUser.AllowUsers {
 			if user.Username == username {
-				sshUsers = append(sshUsers, sshUser)
+				sshUsers[sshUserKey] = sshUser
 			}
 		}
 	}

core/jump/g.go

@@ -38,7 +38,7 @@ func (jps *JumpService) Run(s *ssh.Session) {
 }
 
 func setupConfig() error {
-	fmt.Println("Config file not found. Setup config.", *config.ConfPath)
+	logger.Logger.Info("Config file not found. Setup config.", *config.ConfPath)
 	_, _, err := pui.CreateUser(false, true, nil)
 	if err != nil {
 		return err
@@ -73,7 +73,7 @@ func Configurate() error {
 	if *config.ConfPath == "" {
 		return errors.New("Please specify a config file. ")
 	}
-	fmt.Println("Read config file", *config.ConfPath)
+	logger.Logger.Info("Read config file", *config.ConfPath)
 	if !config.ConfigFileExisted(*config.ConfPath) {
 		err := setupConfig()
 		return err

core/pui/g.go

@@ -1,7 +1,7 @@
 package pui
 
 import (
-	"fmt"
+	"errors"
 
 	"github.com/TNK-Studio/gortal/core/sshd"
 	"github.com/TNK-Studio/gortal/utils/logger"
@@ -20,17 +20,17 @@ func (ui *PUI) SetSession(s *ssh.Session) {
 }
 
 // ShowMenu show menu
-func (ui *PUI) ShowMenu(label string, menu *[]MenuItem, BackOptionLabel string) {
+func (ui *PUI) ShowMenu(label string, menu *[]*MenuItem, BackOptionLabel string, selectedChain []*MenuItem) {
 	for {
 		menuLabels := make([]string, 0)
-		menuItems := make([]MenuItem, 0)
+		menuItems := make([]*MenuItem, 0)
 
 		if menu == nil {
-			return
+			break
 		}
 
 		for index, menuItem := range *menu {
-			if menuItem.IsShow == nil || menuItem.IsShow(index, &menuItem, ui.sess) {
+			if menuItem.IsShow == nil || menuItem.IsShow(index, menuItem, ui.sess, selectedChain) {
 				menuLabels = append(menuLabels, menuItem.Label)
 				menuItems = append(menuItems, menuItem)
 			}
@@ -50,21 +50,22 @@ func (ui *PUI) ShowMenu(label string, menu *[]MenuItem, BackOptionLabel string)
 
 		logger.Logger.Debugf("Selected index: %d subMenuLabel: %+v err: %s", index, subMenuLabel, err)
 		if err != nil {
-			fmt.Printf("Select menu error %s\n", err)
-			return
+			logger.Logger.Infof("Select menu error %s\n", err)
+			break
 		}
 
 		if index == backIndex {
 			break
 		}
-		selected := &(menuItems[index])
+
+		selected := menuItems[index]
 
 		logger.Logger.Debugf("Selected: %+v", selected)
 
 		if selected.GetSubMenu != nil {
 
 			getSubMenu := selected.GetSubMenu
-			subMenu := getSubMenu(index, selected, ui.sess)
+			subMenu := getSubMenu(index, selected, ui.sess, selectedChain)
 
 			if subMenu != nil && len(*subMenu) > 0 {
 				back := "back"
@@ -75,26 +76,33 @@ func (ui *PUI) ShowMenu(label string, menu *[]MenuItem, BackOptionLabel string)
 				if selected.SubMenuTitle != "" {
 					subMenuLabel = selected.SubMenuTitle
 				}
-				ui.ShowMenu(subMenuLabel, subMenu, back)
+				ui.ShowMenu(subMenuLabel, subMenu, back, append(selectedChain, selected))
+			} else {
+				noSubMenuInfo := "No options under this menu ... "
+				if selected.NoSubMenuInfo != "" {
+					noSubMenuInfo = selected.NoSubMenuInfo
+				}
+				sshd.ErrorInfo(errors.New(noSubMenuInfo), ui.sess)
 			}
 		}
 
 		if selected.SelectedFunc != nil {
 			selectedFunc := selected.SelectedFunc
 			logger.Logger.Debugf("Run selectFunc %+v", selectedFunc)
-			err := selectedFunc(index, selected, ui.sess)
+			err := selectedFunc(index, selected, ui.sess, selectedChain)
 			if err != nil {
 				logger.Logger.Errorf("Run selected func err: %s", err)
 				sshd.ErrorInfo(err, ui.sess)
 			}
 			if selected.BackAfterSelected == true {
-				return
+				break
 			}
 		}
 	}
 }
 
 // ShowMainMenu show main menu
 func (ui *PUI) ShowMainMenu() {
-	ui.ShowMenu("Please select the function you need", MainMenu, "Quit")
+	selectedChain := make([]*MenuItem, 0)
+	ui.ShowMenu("Please select the function you need", MainMenu, "Quit", selectedChain)
 }