github/bettercap
1
0
Fork 0
mirror of https://github.com/bettercap/bettercap synced 2025-08-21 05:53:20 -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

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 {