refact: migrated aliases to islazy/data.UnsortedKV

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"
+  revision = "ba851ad172f4be37fcaab106c8c9ebe42e3fa4ac"
-  version = "v1.6.0"
+  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",

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"

network/aliases.go

@@ -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
-}

network/aliases_test.go

@@ -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")
-	}
-}

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, "")
 }

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 {

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

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
+)

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
+		}
+	}
+}