diff --git a/Gopkg.lock b/Gopkg.lock index ade218a1..21804bb1 100644 --- a/Gopkg.lock +++ b/Gopkg.lock @@ -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", diff --git a/Gopkg.toml b/Gopkg.toml index 38c8ed36..87d848e4 100644 --- a/Gopkg.toml +++ b/Gopkg.toml @@ -25,7 +25,7 @@ # unused-packages = true [[constraint]] name = "github.com/evilsocket/islazy" - version = "1.6.0" + version = "1.7.0" [[constraint]] branch = "master" diff --git a/network/aliases.go b/network/aliases.go deleted file mode 100644 index 931e08fe..00000000 --- a/network/aliases.go +++ /dev/null @@ -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 -} diff --git a/network/aliases_test.go b/network/aliases_test.go deleted file mode 100644 index 13da477e..00000000 --- a/network/aliases_test.go +++ /dev/null @@ -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") - } -} diff --git a/network/lan.go b/network/lan.go index d2e56bac..af03b67a 100644 --- a/network/lan.go +++ b/network/lan.go @@ -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, "") } diff --git a/network/net.go b/network/net.go index 9deea45b..3e11a6b7 100644 --- a/network/net.go +++ b/network/net.go @@ -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 { diff --git a/vendor/github.com/evilsocket/islazy/data/doc.go b/vendor/github.com/evilsocket/islazy/data/doc.go new file mode 100644 index 00000000..f0c2351f --- /dev/null +++ b/vendor/github.com/evilsocket/islazy/data/doc.go @@ -0,0 +1,3 @@ +// Package data contains basic threadsafe data structures with +// filesystem persistance and configurable flushing policies. +package data diff --git a/vendor/github.com/evilsocket/islazy/data/flush.go b/vendor/github.com/evilsocket/islazy/data/flush.go new file mode 100644 index 00000000..d2320216 --- /dev/null +++ b/vendor/github.com/evilsocket/islazy/data/flush.go @@ -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 +) diff --git a/vendor/github.com/evilsocket/islazy/data/unsortedkv.go b/vendor/github.com/evilsocket/islazy/data/unsortedkv.go new file mode 100644 index 00000000..ecf32faf --- /dev/null +++ b/vendor/github.com/evilsocket/islazy/data/unsortedkv.go @@ -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 + } + } +}