From 7dccb87aa81c98754a65e10790bd13193d5015a2 Mon Sep 17 00:00:00 2001
From: evilsocket <evilsocket@gmail.com>
Date: Thu, 11 Jan 2018 16:38:02 +0100
Subject: [PATCH] new: custom prompt

---
 session/prompt.go  | 86 ++++++++++++++++++++++++++++++++++++++++++++++
 session/session.go | 13 +++----
 2 files changed, 91 insertions(+), 8 deletions(-)
 create mode 100644 session/prompt.go

diff --git a/session/prompt.go b/session/prompt.go
new file mode 100644
index 00000000..b7bb62c5
--- /dev/null
+++ b/session/prompt.go
@@ -0,0 +1,86 @@
+package session
+
+import (
+	"regexp"
+	"strings"
+
+	"github.com/evilsocket/bettercap-ng/core"
+)
+
+const (
+	PromptVariable = "$"
+	DefaultPrompt  = "{by}{fw}{iface.cidr} {fb}> {iface.addr} {reset} {bold}» {reset}"
+)
+
+var PromptEffects = map[string]string{
+	"{bold}":  core.BOLD,
+	"{dim}":   core.DIM,
+	"{r}":     core.RED,
+	"{g}":     core.GREEN,
+	"{b}":     core.BLUE,
+	"{y}":     core.YELLOW,
+	"{fb}":    core.FG_BLACK,
+	"{fw}":    core.FG_WHITE,
+	"{bdg}":   core.BG_DGRAY,
+	"{br}":    core.BG_RED,
+	"{bg}":    core.BG_GREEN,
+	"{by}":    core.BG_YELLOW,
+	"{blb}":   core.BG_LBLUE, // Ziggy this is for you <3
+	"{reset}": core.RESET,
+}
+
+var PromptCallbacks = map[string]func(s *Session) string{
+	"{iface.cidr}": func(s *Session) string {
+		return s.Interface.CIDR()
+	},
+	"{iface.addr}": func(s *Session) string {
+		return s.Interface.IpAddress
+	},
+	"{iface.mac}": func(s *Session) string {
+		return s.Interface.HwAddress
+	},
+	"{gw.addr}": func(s *Session) string {
+		return s.Gateway.IpAddress
+	},
+	"{gw.mac}": func(s *Session) string {
+		return s.Gateway.HwAddress
+	},
+}
+
+var envRe = regexp.MustCompile("{env\\.(.+)}")
+
+type Prompt struct {
+}
+
+func NewPrompt() Prompt {
+	return Prompt{}
+}
+
+func (p Prompt) Render(s *Session) string {
+	found, prompt := s.Env.Get(PromptVariable)
+	if found == false {
+		prompt = DefaultPrompt
+	}
+
+	for tok, effect := range PromptEffects {
+		prompt = strings.Replace(prompt, tok, effect, -1)
+	}
+
+	for tok, cb := range PromptCallbacks {
+		prompt = strings.Replace(prompt, tok, cb(s), -1)
+	}
+
+	m := envRe.FindStringSubmatch(prompt)
+	if len(m) == 2 {
+		name := m[1]
+		_, value := s.Env.Get(name)
+		prompt = strings.Replace(prompt, m[0], value, -1)
+	}
+
+	// make sure an user error does not screw all terminal
+	if strings.HasPrefix(prompt, core.RESET) == false {
+		prompt += core.RESET
+	}
+
+	return prompt
+}
diff --git a/session/session.go b/session/session.go
index 1b7d3dd2..7a9919a8 100644
--- a/session/session.go
+++ b/session/session.go
@@ -32,6 +32,7 @@ type Session struct {
 	Queue     *packets.Queue           `json:"-"`
 	Input     *readline.Instance       `json:"-"`
 	Active    bool                     `json:"active"`
+	Prompt    Prompt                   `json:"-"`
 
 	CoreHandlers []CommandHandler `json:"-"`
 	Modules      []Module         `json:"-"`
@@ -44,6 +45,7 @@ func New() (*Session, error) {
 	var err error
 
 	s := &Session{
+		Prompt: NewPrompt(),
 		Env:    nil,
 		Active: false,
 		Queue:  nil,
@@ -161,6 +163,8 @@ func (s *Session) Start() error {
 		return err
 	}
 
+	s.Env.Set(PromptVariable, DefaultPrompt)
+
 	s.Env.Set("iface.name", s.Interface.Name())
 	s.Env.Set("iface.address", s.Interface.IpAddress)
 	s.Env.Set("iface.mac", s.Interface.HwAddress)
@@ -224,14 +228,7 @@ func (s *Session) Start() error {
 }
 
 func (s *Session) ReadLine() (string, error) {
-	prompt := core.FG_WHITE + core.BG_YELLOW + " " + s.Interface.CIDR() +
-		core.FG_BLACK +
-		" > " +
-		s.Interface.IpAddress +
-		" " + core.RESET +
-		core.BOLD + " » " + core.RESET
-
-	s.Input.SetPrompt(prompt)
+	s.Input.SetPrompt(s.Prompt.Render(s))
 	s.Input.Refresh()
 	return s.Input.Readline()
 }