mirror of
https://github.com/bettercap/bettercap
synced 2025-08-14 02:36:57 -07:00
refact: refactored caplet loading and command completion logic
This commit is contained in:
evilsocket
parent
bb49de4476
commit
cd14c17ca9
2 changed files with 151 additions and 125 deletions
135
session/caplet.go
Normal file
135
session/caplet.go
Normal file
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue