From abca005651992537bb8885d98c7ec8eeaddf9ee5 Mon Sep 17 00:00:00 2001
From: evilsocket <evilsocket@gmail.com>
Date: Thu, 22 Feb 2018 15:17:36 +0100
Subject: [PATCH] new: implemented 'env' builtin function for proxy modules in
 order to get and set session variables

---
 modules/http_proxy_script_builtins.go | 29 +++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/modules/http_proxy_script_builtins.go b/modules/http_proxy_script_builtins.go
index a39d7a24..86b71105 100644
--- a/modules/http_proxy_script_builtins.go
+++ b/modules/http_proxy_script_builtins.go
@@ -11,6 +11,7 @@ import (
 
 // define functions available to proxy scripts
 func (s *ProxyScript) defineBuiltins() error {
+	// used to read a file ... doh
 	s.VM.Set("readFile", func(call otto.FunctionCall) otto.Value {
 		filename := call.Argument(0).String()
 		raw, err := ioutil.ReadFile(filename)
@@ -27,6 +28,7 @@ func (s *ProxyScript) defineBuiltins() error {
 		return v
 	})
 
+	// log something
 	s.VM.Set("log", func(call otto.FunctionCall) otto.Value {
 		for _, v := range call.ArgumentList {
 			fmt.Printf("%s", v.String())
@@ -37,5 +39,32 @@ func (s *ProxyScript) defineBuiltins() error {
 		return otto.Value{}
 	})
 
+	// read or write environment variable
+	s.VM.Set("env", func(call otto.FunctionCall) otto.Value {
+		argv := call.ArgumentList
+		argc := len(argv)
+
+		if argc == 1 {
+			// get
+			varName := call.Argument(0).String()
+			if found, varValue := s.sess.Env.Get(varName); found == true {
+				v, err := s.VM.ToValue(varValue)
+				if err != nil {
+					log.Error("Could not convert to string: %s", varValue)
+					return otto.Value{}
+				}
+				return v
+			}
+
+		} else if argc == 2 {
+			// set
+			varName := call.Argument(0).String()
+			varValue := call.Argument(1).String()
+			s.sess.Env.Set(varName, varValue)
+		}
+
+		return otto.Value{}
+	})
+
 	return nil
 }