From 20c8844c45cc25ace36976a7093c3c9cce56857a Mon Sep 17 00:00:00 2001
From: Zach Wily <zach@zwily.com>
Date: Wed, 17 Sep 2014 13:06:07 -0600
Subject: [PATCH] support a command to run that specifies the environment for
 the ssh comand

---
 main.go          |  3 ++-
 server/server.go | 15 ++++++++++++++-
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/main.go b/main.go
index 2058346..58eee21 100644
--- a/main.go
+++ b/main.go
@@ -17,6 +17,7 @@ type TunnelDef struct {
 	JumpHost   string `yaml:"jumpHost"`
 	RemoteHost string `yaml:"remoteHost"`
 	RemotePort int    `yaml:"remotePort"`
+	EnvCommand string `yaml:"envCommand"`
 }
 
 type Config struct {
@@ -38,7 +39,7 @@ func main() {
 	goyaml.Unmarshal(configYaml, &config)
 
 	for _, t := range config.Tunnels {
-		s := server.New(t.Name, t.LocalPort, t.JumpHost, t.RemoteHost, t.RemotePort)
+		s := server.New(t.Name, t.LocalPort, t.JumpHost, t.RemoteHost, t.RemotePort, t.EnvCommand)
 		go s.Listen()
 		defer s.Close()
 	}
diff --git a/server/server.go b/server/server.go
index 6725272..4d249aa 100644
--- a/server/server.go
+++ b/server/server.go
@@ -6,6 +6,7 @@ import (
 	"log"
 	"net"
 	"os/exec"
+	"strings"
 	"time"
 )
 
@@ -16,18 +17,20 @@ type Server struct {
 	remoteHost string
 	remotePort int
 	proxyPort  int
+	envCmd     string
 
 	pubsub *pubsub.PubSub
 	cmd    *exec.Cmd
 }
 
-func New(name string, localPort int, jumpHost string, remoteHost string, remotePort int) *Server {
+func New(name string, localPort int, jumpHost string, remoteHost string, remotePort int, envCommand string) *Server {
 	s := &Server{
 		name:       name,
 		localPort:  localPort,
 		jumpHost:   jumpHost,
 		remoteHost: remoteHost,
 		remotePort: remotePort,
+		envCmd:     envCommand,
 	}
 
 	s.pubsub = pubsub.New(0)
@@ -131,6 +134,16 @@ func (s *Server) handlePending(in <-chan *net.TCPConn) {
 					s.jumpHost,
 				)
 
+				if s.envCmd != "" {
+					out, err := exec.Command("/bin/bash", "-c", s.envCmd).Output()
+					if err != nil {
+						log.Fatal(err)
+					}
+
+					log.Printf("%s: using custom environment %s", s.name, out)
+					cmd.Env = strings.Split(string(out), "\n")
+				}
+
 				log.Printf("%s: starting ssh\n", s.name)
 
 				err := cmd.Start()