diff --git a/session/caplet.go b/session/caplet.go new file mode 100644 index 00000000..29619b14 --- /dev/null +++ b/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 +} diff --git a/session/session.go b/session/session.go index f9b0fa10..2f644a29 100644 --- a/session/session.go +++ b/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")), ":") { - if path = core.Trim(path); len(path) > 0 { - CapPaths = append(CapPaths, path) - } - } - - for _, path := range CapPaths { - buildCapletsTree(path, tree) - } - - for root, subElems := range tree { + for root, subElems := range CapletsTree { 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 - 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) + err, caplet := LoadCaplet(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 - } - - 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 { + for _, line := range caplet.Code { + if err = s.Run(line + "\n"); 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 { - return s.runCapletCommand(filename, argv) + if parsed, caplet, argv := parseCapletCommand(line); parsed { + 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) }