diff --git a/core/options.go b/core/options.go index 5344ec86..bf0d1263 100644 --- a/core/options.go +++ b/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`."), diff --git a/session/environment.go b/session/environment.go index bcbd5360..cf1eee4d 100644 --- a/session/environment.go +++ b/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() diff --git a/session/session.go b/session/session.go index e9f9bf26..ff93bbe1 100644 --- a/session/session.go +++ b/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"