mirror of
https://github.com/bettercap/bettercap
synced 2025-08-20 13:33:21 -07:00
refact: migrated aliases to islazy/data.UnsortedKV
This commit is contained in:
parent
2e5b2df7de
commit
7f808ac059
9 changed files with 171 additions and 182 deletions
8
Gopkg.lock
generated
8
Gopkg.lock
generated
|
@ -59,9 +59,10 @@
|
|||
revision = "f58a169a71a51037728990b2d3597a14f56b525b"
|
||||
|
||||
[[projects]]
|
||||
digest = "1:a029ce916ee511044c6b7fc41133249a15e8d3c16219274666205efd6431e9cd"
|
||||
digest = "1:5533d679fd9129a0af86235a01837db24bd66bf368d9a03aaf577a54d3d1e098"
|
||||
name = "github.com/evilsocket/islazy"
|
||||
packages = [
|
||||
"data",
|
||||
"fs",
|
||||
"log",
|
||||
"plugin",
|
||||
|
@ -70,8 +71,8 @@
|
|||
"zip",
|
||||
]
|
||||
pruneopts = "UT"
|
||||
revision = "db3058040a83dba4e35a8931a3e1287c0b802869"
|
||||
version = "v1.6.0"
|
||||
revision = "ba851ad172f4be37fcaab106c8c9ebe42e3fa4ac"
|
||||
version = "v1.7.0"
|
||||
|
||||
[[projects]]
|
||||
branch = "master"
|
||||
|
@ -287,6 +288,7 @@
|
|||
"github.com/chifflier/nfqueue-go/nfqueue",
|
||||
"github.com/dustin/go-humanize",
|
||||
"github.com/elazarl/goproxy",
|
||||
"github.com/evilsocket/islazy/data",
|
||||
"github.com/evilsocket/islazy/fs",
|
||||
"github.com/evilsocket/islazy/log",
|
||||
"github.com/evilsocket/islazy/plugin",
|
||||
|
|
|
@ -25,7 +25,7 @@
|
|||
# unused-packages = true
|
||||
[[constraint]]
|
||||
name = "github.com/evilsocket/islazy"
|
||||
version = "1.6.0"
|
||||
version = "1.7.0"
|
||||
|
||||
[[constraint]]
|
||||
branch = "master"
|
||||
|
|
|
@ -1,99 +0,0 @@
|
|||
package network
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"fmt"
|
||||
"io/ioutil"
|
||||
"os"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"github.com/evilsocket/islazy/fs"
|
||||
"github.com/evilsocket/islazy/str"
|
||||
)
|
||||
|
||||
var fileName, _ = fs.Expand("~/bettercap.aliases")
|
||||
|
||||
type Aliases struct {
|
||||
sync.Mutex
|
||||
|
||||
data map[string]string
|
||||
}
|
||||
|
||||
func LoadAliases() (err error, aliases *Aliases) {
|
||||
aliases = &Aliases{
|
||||
data: make(map[string]string),
|
||||
}
|
||||
|
||||
if fs.Exists(fileName) {
|
||||
var file *os.File
|
||||
|
||||
file, err = os.Open(fileName)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer file.Close()
|
||||
|
||||
scanner := bufio.NewScanner(file)
|
||||
for scanner.Scan() {
|
||||
line := scanner.Text()
|
||||
parts := strings.SplitN(line, " ", 2)
|
||||
mac := str.Trim(parts[0])
|
||||
alias := str.Trim(parts[1])
|
||||
aliases.data[mac] = alias
|
||||
}
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
func (a *Aliases) saveUnlocked() error {
|
||||
data := ""
|
||||
for mac, alias := range a.data {
|
||||
data += fmt.Sprintf("%s %s\n", mac, alias)
|
||||
}
|
||||
return ioutil.WriteFile(fileName, []byte(data), 0644)
|
||||
}
|
||||
|
||||
func (a *Aliases) Save() error {
|
||||
a.Lock()
|
||||
defer a.Unlock()
|
||||
|
||||
return a.saveUnlocked()
|
||||
}
|
||||
|
||||
func (a *Aliases) Get(mac string) string {
|
||||
a.Lock()
|
||||
defer a.Unlock()
|
||||
|
||||
if alias, found := a.data[mac]; found {
|
||||
return alias
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (a *Aliases) Set(mac, alias string) error {
|
||||
a.Lock()
|
||||
defer a.Unlock()
|
||||
|
||||
if alias != "" {
|
||||
a.data[mac] = alias
|
||||
} else {
|
||||
delete(a.data, mac)
|
||||
}
|
||||
|
||||
return a.saveUnlocked()
|
||||
}
|
||||
|
||||
func (a *Aliases) Find(alias string) (mac string, found bool) {
|
||||
a.Lock()
|
||||
defer a.Unlock()
|
||||
|
||||
for m, a := range a.data {
|
||||
if alias == a {
|
||||
return m, true
|
||||
}
|
||||
}
|
||||
|
||||
return "", false
|
||||
}
|
|
@ -1,70 +0,0 @@
|
|||
package network
|
||||
|
||||
import "testing"
|
||||
|
||||
func buildExampleAliases() *Aliases {
|
||||
return &Aliases{}
|
||||
}
|
||||
|
||||
func TestAliasesLoadAliases(t *testing.T) {
|
||||
err, _ := LoadAliases()
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAliasesSaveUnlocked(t *testing.T) {
|
||||
exampleAliases := buildExampleAliases()
|
||||
err := exampleAliases.saveUnlocked()
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAliasesSave(t *testing.T) {
|
||||
exampleAliases := buildExampleAliases()
|
||||
err := exampleAliases.Save()
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAliasesGet(t *testing.T) {
|
||||
exampleAliases := buildExampleAliases()
|
||||
|
||||
exp := ""
|
||||
got := exampleAliases.Get("pi:ca:tw:as:he:re")
|
||||
|
||||
if got != exp {
|
||||
t.Fatalf("expected '%v', got '%v'", exp, got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAliasesSet(t *testing.T) {
|
||||
exampleAliases := buildExampleAliases()
|
||||
exampleAliases.data = make(map[string]string)
|
||||
|
||||
if exampleAliases.Set("pi:ca:tw:as:he:re", "picat") != nil {
|
||||
t.Error("unable to set alias")
|
||||
}
|
||||
|
||||
if exampleAliases.Get("pi:ca:tw:as:he:re") != "picat" {
|
||||
t.Error("unable to get set alias")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAliasesFind(t *testing.T) {
|
||||
exampleAliases := buildExampleAliases()
|
||||
exampleAliases.data = make(map[string]string)
|
||||
err := exampleAliases.Set("pi:ca:tw:as:he:re", "picat")
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
mac, found := exampleAliases.Find("picat")
|
||||
if !found {
|
||||
t.Error("unable to find mac address for alias")
|
||||
}
|
||||
if mac != "pi:ca:tw:as:he:re" {
|
||||
t.Error("unable to find correct mac address for alias")
|
||||
}
|
||||
}
|
|
@ -2,10 +2,12 @@ package network
|
|||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"github.com/evilsocket/islazy/data"
|
||||
"github.com/evilsocket/islazy/fs"
|
||||
)
|
||||
|
||||
const LANDefaultttl = 10
|
||||
|
@ -14,13 +16,15 @@ const LANAliasesFile = "~/bettercap.aliases"
|
|||
type EndpointNewCallback func(e *Endpoint)
|
||||
type EndpointLostCallback func(e *Endpoint)
|
||||
|
||||
var aliasesFileName, _ = fs.Expand(LANAliasesFile)
|
||||
|
||||
type LAN struct {
|
||||
sync.Mutex
|
||||
hosts map[string]*Endpoint
|
||||
iface *Endpoint
|
||||
gateway *Endpoint
|
||||
ttl map[string]uint
|
||||
aliases *Aliases
|
||||
aliases *data.UnsortedKV
|
||||
newCb EndpointNewCallback
|
||||
lostCb EndpointLostCallback
|
||||
}
|
||||
|
@ -30,9 +34,9 @@ type lanJSON struct {
|
|||
}
|
||||
|
||||
func NewLAN(iface, gateway *Endpoint, newcb EndpointNewCallback, lostcb EndpointLostCallback) *LAN {
|
||||
err, aliases := LoadAliases()
|
||||
aliases, err := data.NewUnsortedKV(aliasesFileName, data.FlushOnEdit)
|
||||
if err != nil {
|
||||
fmt.Printf("%s\n", err)
|
||||
panic(err)
|
||||
}
|
||||
|
||||
return &LAN{
|
||||
|
@ -105,7 +109,7 @@ func (lan *LAN) List() (list []*Endpoint) {
|
|||
return
|
||||
}
|
||||
|
||||
func (lan *LAN) Aliases() *Aliases {
|
||||
func (lan *LAN) Aliases() *data.UnsortedKV {
|
||||
return lan.aliases
|
||||
}
|
||||
|
||||
|
@ -197,7 +201,7 @@ func (lan *LAN) AddIfNew(ip, mac string) *Endpoint {
|
|||
return t
|
||||
}
|
||||
|
||||
e := NewEndpointWithAlias(ip, mac, lan.aliases.Get(mac))
|
||||
e := NewEndpointWithAlias(ip, mac, lan.aliases.GetOr(mac, ""))
|
||||
|
||||
lan.hosts[mac] = e
|
||||
lan.ttl[mac] = LANDefaultttl
|
||||
|
@ -208,5 +212,5 @@ func (lan *LAN) AddIfNew(ip, mac string) *Endpoint {
|
|||
}
|
||||
|
||||
func (lan *LAN) GetAlias(mac string) string {
|
||||
return lan.aliases.Get(mac)
|
||||
return lan.aliases.GetOr(mac, "")
|
||||
}
|
||||
|
|
|
@ -7,6 +7,7 @@ import (
|
|||
"regexp"
|
||||
"strings"
|
||||
|
||||
"github.com/evilsocket/islazy/data"
|
||||
"github.com/evilsocket/islazy/str"
|
||||
|
||||
"github.com/malfunkt/iprange"
|
||||
|
@ -67,7 +68,7 @@ func NormalizeMac(mac string) string {
|
|||
return strings.ToLower(strings.Join(parts, ":"))
|
||||
}
|
||||
|
||||
func ParseTargets(targets string, aliasMap *Aliases) (ips []net.IP, macs []net.HardwareAddr, err error) {
|
||||
func ParseTargets(targets string, aliasMap *data.UnsortedKV) (ips []net.IP, macs []net.HardwareAddr, err error) {
|
||||
ips = make([]net.IP, 0)
|
||||
macs = make([]net.HardwareAddr, 0)
|
||||
|
||||
|
@ -90,7 +91,7 @@ func ParseTargets(targets string, aliasMap *Aliases) (ips []net.IP, macs []net.H
|
|||
|
||||
// check and resolve aliases
|
||||
for _, alias := range aliasParser.FindAllString(targets, -1) {
|
||||
if mac, found := aliasMap.Find(alias); found {
|
||||
if mac, found := aliasMap.Get(alias); found {
|
||||
mac = NormalizeMac(mac)
|
||||
hw, err := net.ParseMAC(mac)
|
||||
if err != nil {
|
||||
|
|
3
vendor/github.com/evilsocket/islazy/data/doc.go
generated
vendored
Normal file
3
vendor/github.com/evilsocket/islazy/data/doc.go
generated
vendored
Normal file
|
@ -0,0 +1,3 @@
|
|||
// Package data contains basic threadsafe data structures with
|
||||
// filesystem persistance and configurable flushing policies.
|
||||
package data
|
14
vendor/github.com/evilsocket/islazy/data/flush.go
generated
vendored
Normal file
14
vendor/github.com/evilsocket/islazy/data/flush.go
generated
vendored
Normal file
|
@ -0,0 +1,14 @@
|
|||
package data
|
||||
|
||||
// FlushPolicy is the type of flush policy to use.
|
||||
type FlushPolicy int
|
||||
|
||||
const (
|
||||
// FlushOnEdit saves the object to disk after every modification.
|
||||
FlushOnEdit FlushPolicy = iota
|
||||
// FlushExplicit saves the object to disk only if the Flush method of
|
||||
// the object is explicitly called.
|
||||
FlushExplicit
|
||||
// FlushNone never saves the object to disk.
|
||||
FlushNone
|
||||
)
|
134
vendor/github.com/evilsocket/islazy/data/unsortedkv.go
generated
vendored
Normal file
134
vendor/github.com/evilsocket/islazy/data/unsortedkv.go
generated
vendored
Normal file
|
@ -0,0 +1,134 @@
|
|||
package data
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/gob"
|
||||
"io/ioutil"
|
||||
"os"
|
||||
"sync"
|
||||
|
||||
"github.com/evilsocket/islazy/fs"
|
||||
)
|
||||
|
||||
// UnsortedKV is a thread safe and unsorted key-value
|
||||
// storage with optional persistency on disk.
|
||||
type UnsortedKV struct {
|
||||
sync.Mutex
|
||||
fileName string
|
||||
m map[string]string
|
||||
policy FlushPolicy
|
||||
}
|
||||
|
||||
// NewUnsortedKV creates a new UnsortedKV with the given flush policy.
|
||||
// If fileName already exists, it will be deserialized and loaded.
|
||||
func NewUnsortedKV(fileName string, flushPolicy FlushPolicy) (*UnsortedKV, error) {
|
||||
ukv := &UnsortedKV{
|
||||
fileName: fileName,
|
||||
m: make(map[string]string),
|
||||
policy: flushPolicy,
|
||||
}
|
||||
|
||||
if fileName != "" && fs.Exists(fileName) {
|
||||
raw, err := ioutil.ReadFile(fileName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
decoder := gob.NewDecoder(bytes.NewReader(raw))
|
||||
if err = decoder.Decode(&ukv.m); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
return ukv, nil
|
||||
}
|
||||
|
||||
// Has return true if name exists in the store.
|
||||
func (u *UnsortedKV) Has(name string) bool {
|
||||
u.Lock()
|
||||
defer u.Unlock()
|
||||
_, found := u.m[name]
|
||||
return found
|
||||
}
|
||||
|
||||
// Get return the value of the named object if present, or returns
|
||||
// found as false otherwise.
|
||||
func (u *UnsortedKV) Get(name string) (v string, found bool) {
|
||||
u.Lock()
|
||||
defer u.Unlock()
|
||||
v, found = u.m[name]
|
||||
return
|
||||
}
|
||||
|
||||
// GetOr will return the value of the named object if present,
|
||||
// or a default value.
|
||||
func (u *UnsortedKV) GetOr(name, or string) string {
|
||||
if v, found := u.Get(name); found {
|
||||
return v
|
||||
}
|
||||
return or
|
||||
}
|
||||
|
||||
func (u *UnsortedKV) flushUnlocked() error {
|
||||
buf := new(bytes.Buffer)
|
||||
encoder := gob.NewEncoder(buf)
|
||||
if err := encoder.Encode(u.m); err != nil {
|
||||
return err
|
||||
}
|
||||
return ioutil.WriteFile(u.fileName, buf.Bytes(), os.ModePerm)
|
||||
}
|
||||
|
||||
// Flush flushes the store to disk if the flush policy
|
||||
// is different than FlushNone
|
||||
func (u *UnsortedKV) Flush() error {
|
||||
u.Lock()
|
||||
defer u.Unlock()
|
||||
if u.policy != FlushNone {
|
||||
return u.flushUnlocked()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (u *UnsortedKV) onEdit() error {
|
||||
if u.policy == FlushOnEdit {
|
||||
return u.flushUnlocked()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Set sets a value for a named object.
|
||||
func (u *UnsortedKV) Set(name, value string) error {
|
||||
u.Lock()
|
||||
defer u.Unlock()
|
||||
u.m[name] = value
|
||||
return u.onEdit()
|
||||
}
|
||||
|
||||
// Del deletes a named object from the store.
|
||||
func (u *UnsortedKV) Del(name string) error {
|
||||
u.Lock()
|
||||
defer u.Unlock()
|
||||
delete(u.m, name)
|
||||
return u.onEdit()
|
||||
}
|
||||
|
||||
// Clear deletes every named object from the store.
|
||||
func (u *UnsortedKV) Clear() error {
|
||||
u.Lock()
|
||||
defer u.Unlock()
|
||||
u.m = make(map[string]string)
|
||||
return u.onEdit()
|
||||
}
|
||||
|
||||
// Each iterates each named object in the store by
|
||||
// executing the callback cb on them, if the callback
|
||||
// returns true the iteration is interrupted.
|
||||
func (u *UnsortedKV) Each(cb func(k, v string) bool) {
|
||||
u.Lock()
|
||||
defer u.Unlock()
|
||||
for k, v := range u.m {
|
||||
if stop := cb(k, v); stop {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
Loading…
Add table
Add a link
Reference in a new issue