github/bettercap
1
0
Fork 0
mirror of https://github.com/bettercap/bettercap synced 2025-08-20 13:33:21 -07:00
Code Issues Projects Releases Packages Wiki Activity Actions

refact: migrated aliases to islazy/data.UnsortedKV

Browse source
This commit is contained in:
evilsocket 2018-10-13 18:39:31 +02:00
commit 7f808ac059
No known key found for this signature in database
GPG key ID: 1564D7F30393A456
9 changed files with 171 additions and 182 deletions
Download patch file Download diff file Expand all files Collapse all files

8
Gopkg.lock generated
View file

@ -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",

2
Gopkg.toml
View file

@ -25,7 +25,7 @@
# unused-packages = true
[[constraint]]
name = "github.com/evilsocket/islazy"
version = "1.6.0"
version = "1.7.0"
[[constraint]]
branch = "master"

99
network/aliases.go
View file

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

70
network/aliases_test.go
View file

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

18
network/lan.go
View file

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

5
network/net.go
View file

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