From e216cdde48b3c6f3e20bc53150e3d5b872d97ec9 Mon Sep 17 00:00:00 2001
From: evilsocket <evilsocket@gmail.com>
Date: Sun, 28 Jan 2018 21:07:21 +0100
Subject: [PATCH] new: aliases are now persistant

---
 session/session.go               |  4 +-
 session/session_core_handlers.go |  3 +-
 session/targets.go               | 71 +++++++++++++++++++++++++++++++-
 3 files changed, 74 insertions(+), 4 deletions(-)

diff --git a/session/session.go b/session/session.go
index e311c528..99b918dc 100644
--- a/session/session.go
+++ b/session/session.go
@@ -20,6 +20,8 @@ import (
 	"github.com/evilsocket/bettercap-ng/packets"
 )
 
+const HistoryFile = "~/bettercap.history"
+
 var (
 	I = (*Session)(nil)
 
@@ -144,7 +146,7 @@ func (s *Session) setupInput() error {
 
 	history := ""
 	if *s.Options.NoHistory == false {
-		history = "bettercap.history"
+		history, _ = core.ExpandPath(HistoryFile)
 	}
 
 	cfg := readline.Config{
diff --git a/session/session_core_handlers.go b/session/session_core_handlers.go
index 0dcd9ab3..a49acf90 100644
--- a/session/session_core_handlers.go
+++ b/session/session_core_handlers.go
@@ -183,8 +183,7 @@ func (s *Session) aliasHandler(args []string, sess *Session) error {
 	mac := args[0]
 	alias := args[1]
 
-	if t, found := s.Targets.Targets[mac]; found == true {
-		t.Alias = alias
+	if s.Targets.SetAliasFor(mac, alias) == true {
 		return nil
 	} else {
 		return fmt.Errorf("Could not find endpoint %s", mac)
diff --git a/session/targets.go b/session/targets.go
index 97829cad..a1718f69 100644
--- a/session/targets.go
+++ b/session/targets.go
@@ -1,11 +1,19 @@
 package session
 
 import (
+	"bufio"
+	"fmt"
+	"io/ioutil"
+	"os"
+	"strings"
 	"sync"
 
+	"github.com/evilsocket/bettercap-ng/core"
 	"github.com/evilsocket/bettercap-ng/net"
 )
 
+const TargetsAliasesFile = "~/bettercap.aliases"
+
 type Targets struct {
 	sync.Mutex
 
@@ -14,16 +22,72 @@ type Targets struct {
 	Gateway   *net.Endpoint
 	Targets   map[string]*net.Endpoint
 	TTL       map[string]uint
+	Aliases   map[string]string
+
+	aliasesFileName string
 }
 
 func NewTargets(s *Session, iface, gateway *net.Endpoint) *Targets {
-	return &Targets{
+	t := &Targets{
 		Session:   s,
 		Interface: iface,
 		Gateway:   gateway,
 		Targets:   make(map[string]*net.Endpoint),
 		TTL:       make(map[string]uint),
+		Aliases:   make(map[string]string),
 	}
+
+	t.aliasesFileName, _ = core.ExpandPath(TargetsAliasesFile)
+	if core.Exists(t.aliasesFileName) {
+		if err := t.loadAliases(); err != nil {
+			s.Events.Log(core.ERROR, "%s", err)
+		}
+	}
+
+	return t
+}
+
+func (tp *Targets) loadAliases() error {
+	tp.Session.Events.Log(core.INFO, "Loading aliases from %s ...", tp.aliasesFileName)
+	file, err := os.Open(tp.aliasesFileName)
+	if err != nil {
+		return err
+	}
+	defer file.Close()
+
+	scanner := bufio.NewScanner(file)
+	for scanner.Scan() {
+		line := scanner.Text()
+		parts := strings.SplitN(line, " ", 2)
+		mac := strings.Trim(parts[0], "\r\n\t ")
+		alias := strings.Trim(parts[1], "\r\n\t ")
+		tp.Session.Events.Log(core.DEBUG, " aliases[%s] = '%s'", mac, alias)
+		tp.Aliases[mac] = alias
+	}
+
+	return nil
+}
+
+func (tp *Targets) saveAliases() {
+	data := ""
+	for mac, alias := range tp.Aliases {
+		data += fmt.Sprintf("%s %s\n", mac, alias)
+	}
+	ioutil.WriteFile(tp.aliasesFileName, []byte(data), 0644)
+}
+
+func (tp *Targets) SetAliasFor(mac, alias string) bool {
+	tp.Lock()
+	defer tp.Unlock()
+
+	if t, found := tp.Targets[mac]; found == true {
+		tp.Aliases[mac] = alias
+		t.Alias = alias
+		tp.saveAliases()
+		return true
+	}
+
+	return false
 }
 
 func (tp *Targets) Remove(ip, mac string) {
@@ -66,6 +130,7 @@ func (tp *Targets) AddIfNotExist(ip, mac string) *net.Endpoint {
 		return nil
 	}
 
+	mac = net.NormalizeMac(mac)
 	if t, found := tp.Targets[mac]; found {
 		return t
 	}
@@ -75,6 +140,10 @@ func (tp *Targets) AddIfNotExist(ip, mac string) *net.Endpoint {
 		tp.Session.Events.Add("target.resolved", e)
 	}
 
+	if alias, found := tp.Aliases[mac]; found {
+		e.Alias = alias
+	}
+
 	tp.Targets[mac] = e
 	tp.TTL[mac] = 2