WIP: session and recovery

Author: Nicolas Chatelain

cmd/proxy/app/app.go

@@ -27,6 +27,12 @@ var ListenerList map[int]controller.Listener
 var ListenerListMutex sync.Mutex
 var ProxyController *controller.Controller
 
+// CurrentAgentID points to the selected agent in the UI (when running session)
+var CurrentAgentID int
+
+// Store AgentIDs
+var AgentCounter int
+
 var (
 	ErrInvalidAgent   = errors.New("please, select an agent using the session command")
 	ErrAlreadyRunning = errors.New("already running")
@@ -35,21 +41,76 @@ var (
 
 func RegisterAgent(agent *controller.LigoloAgent) error {
 	AgentListMutex.Lock()
-	AgentList[agent.Id] = agent
-	AgentListMutex.Unlock()
+	defer AgentListMutex.Unlock()
+
+	for _, registeredAgents := range AgentList {
+		if agent.SessionID == registeredAgents.SessionID {
+			logrus.Infof("Recovered an agent: %s", registeredAgents.Name)
+			registeredAgents.Session = agent.Session
+			if registeredAgents.Running {
+				go StartTunnel(registeredAgents, registeredAgents.Interface)
+			}
+			return nil
+		}
+	}
+	AgentCounter++
+	AgentList[AgentCounter] = agent
 	return nil
 }
 
 func UnregisterAgent(agent *controller.LigoloAgent) error {
-	AgentListMutex.Lock()
+	/*AgentListMutex.Lock()
 	delete(AgentList, agent.Id)
-	AgentListMutex.Unlock()
+	AgentListMutex.Unlock()*/
 	return nil
 }
 
+func StartTunnel(agent *controller.LigoloAgent, tunName string) {
+	logrus.Infof("Starting tunnel to %s", agent.Name)
+	ligoloStack, err := proxy.NewLigoloTunnel(netstack.StackSettings{
+		TunName:     tunName,
+		MaxInflight: 4096,
+	})
+	if err != nil {
+		logrus.Error("Unable to create tunnel, err:", err)
+		return
+	}
+	ifName, err := ligoloStack.GetStack().Interface().Name()
+	if err != nil {
+		logrus.Warn("unable to get interface name, err:", err)
+		ifName = tunName
+	}
+	agent.Interface = ifName
+	agent.Running = true
+
+	ctx, cancelTunnel := context.WithCancel(context.Background())
+	go ligoloStack.HandleSession(agent.Session, ctx)
+
+	for {
+		select {
+		case <-agent.CloseChan: // User stopped
+			logrus.Infof("Closing tunnel to %s...", agent.Name)
+			cancelTunnel()
+			return
+		case <-agent.Session.CloseChan(): // Agent closed
+			logrus.Warnf("Lost tunnel connection with agent %s!", agent.Name)
+			//agent.Running = false
+			//agent.Session = nil
+
+			if currentAgent, ok := AgentList[CurrentAgentID]; ok {
+				if currentAgent.SessionID == agent.SessionID {
+					App.SetDefaultPrompt()
+					agent.Session = nil
+				}
+			}
+
+			cancelTunnel()
+			return
+		}
+	}
+}
+
 func Run() {
-	// CurrentAgent points to the selected agent in the UI (when running session)
-	var CurrentAgentID int
 	// AgentList contains all the connected agents
 	AgentList = make(map[int]*controller.LigoloAgent)
 	// ListenerList contains all listener relays
@@ -61,7 +122,13 @@ func Run() {
 		Usage: "session",
 		Run: func(c *grumble.Context) error {
 			AgentListMutex.Lock()
-			if len(AgentList) == 0 {
+			sessionCount := 0
+			for _, agent := range AgentList {
+				if agent.Session != nil && !agent.Session.IsClosed() {
+					sessionCount += 1
+				}
+			}
+			if sessionCount == 0 {
 				AgentListMutex.Unlock()
 				return errors.New("no sessions available")
 			}
@@ -72,7 +139,9 @@ func Run() {
 				Options: func() (out []string) {
 					AgentListMutex.Lock()
 					for id, agent := range AgentList {
-						out = append(out, fmt.Sprintf("%d - %s", id, agent.String()))
+						if agent.Session != nil && !agent.Session.IsClosed() {
+							out = append(out, fmt.Sprintf("%d - %s", id, agent.String()))
+						}
 					}
 					AgentListMutex.Unlock()
 					return
@@ -194,48 +263,8 @@ func Run() {
 				}
 			}
 
-			go func() {
-				logrus.Infof("Starting tunnel to %s", CurrentAgent.Name)
-				ligoloStack, err := proxy.NewLigoloTunnel(netstack.StackSettings{
-					TunName:     c.Flags.String("tun"),
-					MaxInflight: 4096,
-				})
-				if err != nil {
-					logrus.Error("Unable to create tunnel, err:", err)
-					return
-				}
-				ifName, err := ligoloStack.GetStack().Interface().Name()
-				if err != nil {
-					logrus.Warn("unable to get interface name, err:", err)
-					ifName = c.Flags.String("tun")
-				}
-				CurrentAgent.Interface = ifName
-				CurrentAgent.Running = true
-
-				ctx, cancelTunnel := context.WithCancel(context.Background())
-				go ligoloStack.HandleSession(CurrentAgent.Session, ctx)
+			go StartTunnel(CurrentAgent, c.Flags.String("tun"))
 
-				for {
-					select {
-					case <-CurrentAgent.CloseChan: // User stopped
-						logrus.Infof("Closing tunnel to %s...", CurrentAgent.Name)
-						cancelTunnel()
-						return
-					case <-CurrentAgent.Session.CloseChan(): // Agent closed
-						logrus.Warnf("Lost connection with agent %s!", CurrentAgent.Name)
-						// Connection lost, we need to delete the Agent from the list
-						AgentListMutex.Lock()
-						delete(AgentList, CurrentAgent.Id)
-						AgentListMutex.Unlock()
-						if CurrentAgent.Id == CurrentAgent.Id {
-							App.SetDefaultPrompt()
-							CurrentAgent.Session = nil
-						}
-						cancelTunnel()
-						return
-					}
-				}
-			}()
 			return nil
 		},
 	})
@@ -252,10 +281,10 @@ func Run() {
 
 			AgentListMutex.Lock()
 
-			for _, agent := range AgentList {
+			for id, agent := range AgentList {
 
 				if agent.Running {
-					t.AppendRow(table.Row{agent.Id, agent.Name, agent.Interface})
+					t.AppendRow(table.Row{id, agent.Name, agent.Interface})
 				}
 			}
 			AgentListMutex.Unlock()

cmd/proxy/main.go

@@ -120,9 +120,6 @@ func main() {
 					select {
 					case <-agent.Session.CloseChan(): // Agent closed
 						logrus.Warnf("Lost ligolo-ng connection with agent %s!", agent.Name)
-						if err := app.UnregisterAgent(agent); err != nil {
-							logrus.Errorf("could not unregister agent: %s", err.Error())
-						}
 						return
 					}
 				}

go.mod

@@ -16,6 +16,7 @@ require (
 )
 
 require (
+	github.com/google/uuid v1.3.0
 	github.com/nicocha30/gvisor-ligolo v0.0.0-20230726075806-989fa2c0a413
 	github.com/vishvananda/netlink v1.1.1-0.20211118161826-650dca95af54
 	golang.org/x/sys v0.21.0
@@ -29,7 +30,6 @@ require (
 	github.com/desertbit/readline v1.5.1 // indirect
 	github.com/fatih/color v1.13.0 // indirect
 	github.com/google/btree v1.1.2 // indirect
-	github.com/google/uuid v1.3.0 // indirect
 	github.com/hashicorp/errwrap v1.1.0 // indirect
 	github.com/hashicorp/go-multierror v1.1.1 // indirect
 	github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51 // indirect

pkg/agent/handler.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"github.com/google/uuid"
 	"net"
 	"os"
 	"os/user"
@@ -21,10 +22,13 @@ var listenerConntrack map[int32]net.Conn
 var listenerMap map[int32]interface{}
 var connTrackID int32
 var listenerID int32
+var sessionID string
 
 func init() {
 	listenerConntrack = make(map[int32]net.Conn)
 	listenerMap = make(map[int32]interface{})
+	id := uuid.New()
+	sessionID = id.String()
 }
 
 // Listener is the base class implementing listener sockets for Ligolo
@@ -171,6 +175,7 @@ func HandleConn(conn net.Conn) {
 		infoResponse := protocol.InfoReplyPacket{
 			Name:       fmt.Sprintf("%s@%s", username, hostname),
 			Interfaces: protocol.NewNetInterfaces(netifaces),
+			SessionID:  sessionID,
 		}
 
 		if err := encoder.Encode(protocol.Envelope{

pkg/controller/agent.go

@@ -7,14 +7,13 @@ import (
 	"net"
 )
 
-var AgentCounter = 0
 var ListenerCounter = 0
 
 type LigoloAgent struct {
-	Id        int
 	Name      string
 	Network   []protocol.NetInterface
 	Session   *yamux.Session
+	SessionID string
 	CloseChan chan bool
 	Interface string
 	Running   bool
@@ -31,11 +30,16 @@ type Listener struct {
 }
 
 func (l Listener) String() string {
-	return fmt.Sprintf("#%d [%s] (%s) [Agent] %s => [Proxy] %s", l.Agent.Id, l.Agent.Name, l.Network, l.ListenerAddr, l.RedirectAddr)
+	return fmt.Sprintf("[%s] (%s) [Agent] %s => [Proxy] %s", l.Agent.Name, l.Network, l.ListenerAddr, l.RedirectAddr)
 }
 
 func (la *LigoloAgent) String() string {
-	return fmt.Sprintf("#%d - %s - %s", la.Id, la.Name, la.Session.RemoteAddr())
+	raddr := "Disconnected"
+	if la.Session != nil {
+		raddr = la.Session.RemoteAddr().String()
+	}
+
+	return fmt.Sprintf("%s - %s - %s", la.Name, raddr, la.SessionID)
 }
 
 func NewAgent(session *yamux.Session) (*LigoloAgent, error) {
@@ -62,12 +66,12 @@ func NewAgent(session *yamux.Session) (*LigoloAgent, error) {
 
 	response := protocolDecoder.Envelope.Payload
 	reply := response.(protocol.InfoReplyPacket)
-	AgentCounter++
+
 	return &LigoloAgent{
-		Id:        AgentCounter,
 		Name:      reply.Name,
 		Network:   reply.Interfaces,
 		Session:   session,
+		SessionID: reply.SessionID,
 		CloseChan: make(chan bool),
 	}, nil
 }

pkg/protocol/packets.go

@@ -47,6 +47,7 @@ type InfoRequestPacket struct {
 type InfoReplyPacket struct {
 	Name       string
 	Interfaces []NetInterface
+	SessionID  string
 }
 
 // ListenerSockRequestPacket is used by the proxy when relaying a listener socket