refact: refactored caplet loading and command completion logic

session/caplet.go

@@ -0,0 +1,135 @@
+package session
+
+import (
+	"bufio"
+	"fmt"
+	"io/ioutil"
+	"os"
+	"path/filepath"
+	"strings"
+	"sync"
+
+	"github.com/bettercap/bettercap/core"
+)
+
+const (
+	CapletSuffix = ".cap"
+)
+
+type Caplet struct {
+	Path string
+	Code []string
+}
+
+var (
+	CapletsTree     = make(map[string][]string)
+	CapletLoadPaths = []string{
+		"./caplets/",
+		"/usr/share/bettercap/caplets/",
+	}
+
+	cache     = make(map[string]*Caplet)
+	cacheLock = sync.Mutex{}
+)
+
+func buildCapletsTree(path string, prefix string) {
+	files, _ := ioutil.ReadDir(path)
+	for _, file := range files {
+		filename := file.Name()
+		if strings.HasSuffix(filename, CapletSuffix) {
+			base := strings.TrimPrefix(path, prefix)
+			name := strings.Replace(filename, CapletSuffix, "", -1)
+			CapletsTree[base+name] = []string{}
+		} else if file.IsDir() {
+			buildCapletsTree(filepath.Join(path, filename)+"/", prefix)
+		}
+	}
+}
+
+func init() {
+	for _, path := range core.SepSplit(core.Trim(os.Getenv("CAPSPATH")), ":") {
+		if path = core.Trim(path); len(path) > 0 {
+			CapletLoadPaths = append(CapletLoadPaths, path)
+		}
+	}
+
+	for i, path := range CapletLoadPaths {
+		CapletLoadPaths[i], _ = filepath.Abs(path)
+	}
+
+	for _, path := range CapletLoadPaths {
+		buildCapletsTree(path, path)
+	}
+}
+
+func LoadCaplet(name string) (error, *Caplet) {
+	cacheLock.Lock()
+	defer cacheLock.Unlock()
+
+	if caplet, found := cache[name]; found {
+		return nil, caplet
+	}
+
+	names := []string{name}
+	if !strings.HasSuffix(name, CapletSuffix) {
+		names = append(names, name+CapletSuffix)
+	}
+
+	for _, path := range CapletLoadPaths {
+		if !strings.HasSuffix(name, CapletSuffix) {
+			name += CapletSuffix
+		}
+		names = append(names, filepath.Join(path, name))
+	}
+
+	for _, filename := range names {
+		if core.Exists(filename) {
+			cap := &Caplet{
+				Path: filename,
+				Code: make([]string, 0),
+			}
+
+			I.Events.Log(core.INFO, "reading from caplet %s ...", filename)
+			input, err := os.Open(filename)
+			if err != nil {
+				return fmt.Errorf("error reading caplet %s: %v", filename, err), nil
+			}
+			defer input.Close()
+
+			scanner := bufio.NewScanner(input)
+			scanner.Split(bufio.ScanLines)
+			for scanner.Scan() {
+				line := core.Trim(scanner.Text())
+				if line == "" || line[0] == '#' {
+					continue
+				}
+				cap.Code = append(cap.Code, line)
+			}
+
+			cache[name] = cap
+			return nil, cap
+		}
+	}
+
+	return fmt.Errorf("caplet %s not found", name), nil
+}
+
+func parseCapletCommand(line string) (is bool, caplet *Caplet, argv []string) {
+	file := core.Trim(line)
+	parts := strings.Split(file, " ")
+	argc := len(parts)
+	argv = make([]string, 0)
+	// check for any arguments
+	if argc > 1 {
+		file = core.Trim(parts[0])
+		if argc >= 2 {
+			argv = parts[1:]
+		}
+	}
+
+	if err, cap := LoadCaplet(file); err == nil {
+		return true, cap, argv
+	}
+
+	return false, nil, nil
+}

session/session.go

@@ -1,13 +1,11 @@
 package session
 
 import (
-	"bufio"
 	"errors"
 	"fmt"
 	"net"
 	"os"
 	"os/signal"
-	"path/filepath"
 	"regexp"
 	"runtime"
 	"runtime/pprof"
@@ -24,7 +22,6 @@ import (
 	"github.com/bettercap/bettercap/packets"
 
 	"github.com/adrianmo/go-nmea"
-	"io/ioutil"
 )
 
 const (
@@ -40,11 +37,6 @@ var (
 
 	reCmdSpaceCleaner = regexp.MustCompile(`^([^\s]+)\s+(.+)$`)
 	reEnvVarCapture   = regexp.MustCompile(`{env\.([^}]+)}`)
-
-	CapPaths = []string{
-		"./caplets/",
-		"/usr/share/bettercap/caplets/",
-	}
 )
 
 type UnknownCommandCallback func(cmd string) bool
@@ -220,24 +212,12 @@ func (s *Session) setupReadline() error {
 		}
 	}
 
-	for _, path := range core.SepSplit(core.Trim(os.Getenv("CAPSPATH")), ":") {
+	for root, subElems := range CapletsTree {
-		if path = core.Trim(path); len(path) > 0 {
-			CapPaths = append(CapPaths, path)
-		}
-	}
-
-	for _, path := range CapPaths {
-		buildCapletsTree(path, tree)
-	}
-
-	for root, subElems := range tree {
 		item := readline.PcItem(root)
 		item.Children = []readline.PrefixCompleterInterface{}
-
 		for _, child := range subElems {
 			item.Children = append(item.Children, readline.PcItem(child))
 		}
-
 		pcompleters = append(pcompleters, item)
 	}
 
@@ -257,21 +237,6 @@ func (s *Session) setupReadline() error {
 	return err
 }
 
-func buildCapletsTree(path string, tree map[string][]string) {
-	_buildCapletsTree(path, tree, path)
-}
-
-func _buildCapletsTree(path string, tree map[string][]string, prefix string) {
-	subFiles, _ := ioutil.ReadDir(path)
-	for _, subF := range subFiles {
-		if strings.HasSuffix(subF.Name(), ".cap") {
-			tree[strings.TrimPrefix(path, prefix)+strings.Replace(subF.Name(), ".cap", "", -1)] = []string{}
-		} else if subF.IsDir() {
-			_buildCapletsTree(path+subF.Name()+"/", tree, prefix)
-		}
-	}
-}
-
 func containsCapitals(s string) bool {
 	for _, ch := range s {
 		if ch < 133 && ch > 101 {
@@ -504,97 +469,14 @@ func (s *Session) ReadLine() (string, error) {
 	return s.Input.Readline()
 }
 
-func (s *Session) getCapletFilePath(caplet string) string {
-	if core.Exists(caplet) {
-		return caplet
-	}
-
-	for _, path := range CapPaths {
-		filename := filepath.Join(path, caplet)
-		if !strings.HasSuffix(filename, ".cap") {
-			filename += ".cap"
-		}
-		if core.Exists(filename) {
-			return filename
-		}
-	}
-
-	return ""
-}
-
 func (s *Session) RunCaplet(filename string) error {
-	var caplet string
+	err, caplet := LoadCaplet(filename)
-	if caplet = s.getCapletFilePath(filename); caplet == "" {
-		return fmt.Errorf("could not load caplet from %s", filename)
-	}
-
-	s.Events.Log(core.INFO, "Reading from caplet %s ...", caplet)
-
-	input, err := os.Open(caplet)
 	if err != nil {
 		return err
 	}
-	defer input.Close()
 
-	scanner := bufio.NewScanner(input)
+	for _, line := range caplet.Code {
-	scanner.Split(bufio.ScanLines)
+		if err = s.Run(line + "\n"); err != nil {
-
-	for scanner.Scan() {
-		line := scanner.Text()
-		if line == "" || line[0] == '#' {
-			continue
-		}
-
-		if err = s.Run(line); err != nil {
-			return err
-		}
-	}
-
-	return nil
-}
-
-func (s *Session) isCapletCommand(line string) (is bool, filename string, argv []string) {
-	file := core.Trim(line)
-	parts := strings.Split(file, " ")
-	argc := len(parts)
-	argv = make([]string, 0)
-	// check for any arguments
-	if argc > 1 {
-		file = core.Trim(parts[0])
-		if argc >= 2 {
-			argv = parts[1:]
-		}
-	}
-
-	if filename := s.getCapletFilePath(file); filename != "" {
-		return true, filename, argv
-	}
-
-	return false, "", nil
-}
-
-func (s *Session) runCapletCommand(filename string, argv []string) error {
-	input, err := os.Open(filename)
-	if err != nil {
-		return err
-	}
-	defer input.Close()
-
-	scanner := bufio.NewScanner(input)
-	scanner.Split(bufio.ScanLines)
-
-	for scanner.Scan() {
-		line := scanner.Text()
-		if line == "" || line[0] == '#' {
-			continue
-		}
-
-		// replace $0 with argv[0], $1 with argv[1] and so on
-		for i, arg := range argv {
-			line = strings.Replace(line, fmt.Sprintf("$%d", i), arg, -1)
-		}
-
-		if err = s.Run(line); err != nil {
 			return err
 		}
 	}
@@ -631,8 +513,17 @@ func (s *Session) Run(line string) error {
 	}
 
 	// is it a caplet command?
-	if is, filename, argv := s.isCapletCommand(line); is {
+	if parsed, caplet, argv := parseCapletCommand(line); parsed {
-		return s.runCapletCommand(filename, argv)
+		for _, line := range caplet.Code {
+			// replace $0 with argv[0], $1 with argv[1] and so on
+			for i, arg := range argv {
+				line = strings.Replace(line, fmt.Sprintf("$%d", i), arg, -1)
+			}
+
+			if err = s.Run(line + "\n"); err != nil {
+				return err
+			}
+		}
 	}
 
 	// is it a proxy module custom command?
@@ -640,5 +531,5 @@ func (s *Session) Run(line string) error {
 		return nil
 	}
 
-	return fmt.Errorf("Unknown or invalid syntax \"%s%s%s\", type %shelp%s for the help menu.", core.BOLD, line, core.RESET, core.BOLD, core.RESET)
+	return fmt.Errorf("unknown or invalid syntax \"%s%s%s\", type %shelp%s for the help menu.", core.BOLD, line, core.RESET, core.BOLD, core.RESET)
 }