From a2cc2746b4f60eb9406aad42f66464f403f6c6c7 Mon Sep 17 00:00:00 2001
From: evilsocket <evilsocket@gmail.com>
Date: Sat, 27 Jan 2018 20:53:09 +0100
Subject: [PATCH] new: new alias command

---
 modules/net_recon_show.go        |  9 +++++++--
 net/endpoint.go                  |  1 +
 session/session_core_handlers.go | 28 ++++++++++++++++++++++++++++
 3 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/modules/net_recon_show.go b/modules/net_recon_show.go
index 16b0b46a..beff2c1f 100644
--- a/modules/net_recon_show.go
+++ b/modules/net_recon_show.go
@@ -107,10 +107,15 @@ func (d *Discovery) Show(by string) error {
 				seen = core.Dim(seen)
 			}
 
+			name := core.Yellow(t.Hostname)
+			if t.Alias != "" {
+				name = core.Green(t.Alias)
+			}
+
 			data[i] = []string{
 				t.IpAddress,
 				t.HwAddress,
-				core.Yellow(t.Hostname),
+				name,
 				t.Vendor,
 				humanize.Bytes(traffic.Sent),
 				humanize.Bytes(traffic.Received),
@@ -120,7 +125,7 @@ func (d *Discovery) Show(by string) error {
 
 		table = tablewriter.NewWriter(os.Stdout)
 
-		table.SetHeader([]string{"IP", "MAC", "Hostname", "Vendor", "Sent", "Recvd", "Last Seen"})
+		table.SetHeader([]string{"IP", "MAC", "Name", "Vendor", "Sent", "Recvd", "Last Seen"})
 		table.SetColWidth(80)
 		table.AppendBulk(data)
 		table.Render()
diff --git a/net/endpoint.go b/net/endpoint.go
index f2a7c776..5690951b 100644
--- a/net/endpoint.go
+++ b/net/endpoint.go
@@ -22,6 +22,7 @@ type Endpoint struct {
 	IpAddressUint32  uint32                 `json:"-"`
 	HwAddress        string                 `json:"mac"`
 	Hostname         string                 `json:"hostname"`
+	Alias            string                 `json:"alias"`
 	Vendor           string                 `json:"vendor"`
 	ResolvedCallback OnHostResolvedCallback `json:"-"`
 	FirstSeen        time.Time              `json:"first_seen"`
diff --git a/session/session_core_handlers.go b/session/session_core_handlers.go
index c0692d18..0dcd9ab3 100644
--- a/session/session_core_handlers.go
+++ b/session/session_core_handlers.go
@@ -179,6 +179,18 @@ func (s *Session) shHandler(args []string, sess *Session) error {
 	return err
 }
 
+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
+		return nil
+	} else {
+		return fmt.Errorf("Could not find endpoint %s", mac)
+	}
+}
+
 func (s *Session) addHandler(h CommandHandler, c *readline.PrefixCompleter) {
 	h.Completer = c
 	s.CoreHandlers = append(s.CoreHandlers, h)
@@ -274,4 +286,20 @@ func (s *Session) registerCoreHandlers() {
 		"Execute a shell command and print its output.",
 		s.shHandler),
 		readline.PcItem("!"))
+
+	s.addHandler(NewCommandHandler("alias MAC NAME",
+		"^alias\\s+([a-fA-F0-9:]{17})\\s+(.+)",
+		"Assign an alias to a given endpoint given its MAC address.",
+		s.aliasHandler),
+		readline.PcItem("alias", readline.PcItemDynamic(func(prefix string) []string {
+			prefix = strings.Trim(prefix[5:], "\t\r\n ")
+			macs := []string{""}
+			for mac, _ := range s.Targets.Targets {
+				if prefix == "" || strings.HasPrefix(mac, prefix) == true {
+					macs = append(macs, mac)
+				}
+			}
+			return macs
+		})))
+
 }