new: new -env-file for env persistance (closes #81)

core/options.go

@@ -8,6 +8,7 @@ type Options struct {
 	Debug         *bool
 	Silent        *bool
 	NoHistory     *bool
+	EnvFile       *string
 	Commands      *string
 	CpuProfile    *string
 	MemProfile    *string
@@ -20,6 +21,7 @@ func ParseOptions() (Options, error) {
 		Debug:         flag.Bool("debug", false, "Print debug messages."),
 		Silent:        flag.Bool("silent", false, "Suppress all logs which are not errors."),
 		NoHistory:     flag.Bool("no-history", false, "Disable interactive session history file."),
+		EnvFile:       flag.String("env-file", "~/bettercap.env", "Load environment variables from this file if found, set to empty to disable environment persistance."),
 		Commands:      flag.String("eval", "", "Run one or more commands separated by ; in the interactive session, used to set variables via command line."),
 		CpuProfile:    flag.String("cpu-profile", "", "Write cpu profile `file`."),
 		MemProfile:    flag.String("mem-profile", "", "Write memory profile to `file`."),

session/environment.go

@@ -1,7 +1,9 @@
 package session
 
 import (
+	"encoding/json"
 	"fmt"
+	"io/ioutil"
 	"sort"
 	"strconv"
 	"sync"
@@ -21,7 +23,7 @@ type Environment struct {
 	sess *Session
 }
 
-func NewEnvironment(s *Session) *Environment {
+func NewEnvironment(s *Session, envFile string) *Environment {
 	env := &Environment{
 		Padding: 0,
 		Data:    make(map[string]string),
@@ -29,9 +31,42 @@ func NewEnvironment(s *Session) *Environment {
 		cbs:     make(map[string]SetCallback),
 	}
 
+	if envFile != "" {
+		envFile, _ := core.ExpandPath(envFile)
+		if core.Exists(envFile) {
+			if err := env.Load(envFile); err != nil {
+				fmt.Printf("Error while loading %s: %s\n", envFile, err)
+			}
+		}
+	}
+
 	return env
 }
 
+func (env *Environment) Load(fileName string) error {
+	env.Lock()
+	defer env.Unlock()
+
+	raw, err := ioutil.ReadFile(fileName)
+	if err != nil {
+		return err
+	}
+
+	return json.Unmarshal(raw, &env.Data)
+}
+
+func (env *Environment) Save(fileName string) error {
+	env.Lock()
+	defer env.Unlock()
+
+	raw, err := json.Marshal(env.Data)
+	if err != nil {
+		return err
+	}
+
+	return ioutil.WriteFile(fileName, raw, 0644)
+}
+
 func (env *Environment) Has(name string) bool {
 	env.Lock()
 	defer env.Unlock()

session/session.go

@@ -142,7 +142,7 @@ func New() (*Session, error) {
 		}
 	}
 
-	s.Env = NewEnvironment(s)
+	s.Env = NewEnvironment(s, *s.Options.EnvFile)
 	s.Events = NewEventPool(*s.Options.Debug, *s.Options.Silent)
 
 	s.registerCoreHandlers()
@@ -246,6 +246,13 @@ func (s *Session) Close() {
 	s.Firewall.Restore()
 	s.Queue.Stop()
 
+	if *s.Options.EnvFile != "" {
+		envFile, _ := core.ExpandPath(*s.Options.EnvFile)
+		if err := s.Env.Save(envFile); err != nil {
+			fmt.Printf("Error while storing the environment to %s: %s", envFile, err)
+		}
+	}
+
 	if *s.Options.CpuProfile != "" {
 		pprof.StopCPUProfile()
 	}
@@ -304,7 +311,6 @@ func (s *Session) setupSignals() {
 }
 
 func (s *Session) setupEnv() {
-	s.Env.Set(PromptVariable, DefaultPrompt)
 	s.Env.Set("iface.index", fmt.Sprintf("%d", s.Interface.Index))
 	s.Env.Set("iface.name", s.Interface.Name())
 	s.Env.Set("iface.ipv4", s.Interface.IpAddress)
@@ -313,6 +319,10 @@ func (s *Session) setupEnv() {
 	s.Env.Set("gateway.address", s.Gateway.IpAddress)
 	s.Env.Set("gateway.mac", s.Gateway.HwAddress)
 
+	if found, v := s.Env.Get(PromptVariable); found == false || v == "" {
+		s.Env.Set(PromptVariable, DefaultPrompt)
+	}
+
 	dbg := "false"
 	if *s.Options.Debug {
 		dbg = "true"