From 915b429eefba51f35532d22c4e2177ced6207d8f Mon Sep 17 00:00:00 2001
From: lobsterjerusalem <work@nday.dev>
Date: Fri, 9 May 2025 11:05:28 -0600
Subject: [PATCH] Added a LastSeen value to session and added automatic
 updating of it. Also added an Active param to the Session struct and changed
 the logic to RemoveSession(s) functions to set this value to false instead of
 deleting it

---
 c2/channel/channel.go | 64 +++++++++++++++++++++++++++++++++++++------
 c2/cli/basic.go       |  7 +++++
 2 files changed, 63 insertions(+), 8 deletions(-)

diff --git a/c2/channel/channel.go b/c2/channel/channel.go
index 3e51a33..e22b324 100644
--- a/c2/channel/channel.go
+++ b/c2/channel/channel.go
@@ -31,6 +31,8 @@ type Session struct {
 	RemoteAddr     string
 	ConnectionTime time.Time
 	conn           *net.Conn
+	Active         bool
+	LastSeen       time.Time
 }
 
 // HasSessions checks if a channel has any tracked sessions. This can be used to lookup if a C2
@@ -39,7 +41,13 @@ type Session struct {
 //	c, ok := c2.GetInstance(conf.C2Type)
 //	c.Channel().HasSessions()
 func (c *Channel) HasSessions() bool {
-	return len(c.Sessions) > 0
+	for _, sess := range c.Sessions {
+		if sess.Active {
+			return true
+		}
+	}
+
+	return false
 }
 
 // AddSession adds a remote connection for session tracking. If a network connection is being
@@ -65,11 +73,53 @@ func (c *Channel) AddSession(conn *net.Conn, addr string) bool {
 		ConnectionTime: time.Now(),
 		conn:           conn,
 		RemoteAddr:     addr,
+		LastSeen:       time.Now(),
+		Active:         true,
+	}
+
+	return true
+}
+
+// Updates the LastSeen value for provided connection to the provided time
+func (c *Channel) UpdateLastSeenByConn(conn net.Conn, timeStamp time.Time) bool {
+	id, ok := c.GetSessionIDByConn(conn)
+	if !ok {
+		return false
+	}
+
+	session, ok := c.Sessions[id]
+	if !ok {
+		output.PrintFrameworkError("Session ID does not exist")
+
+		return false
 	}
 
+	session.LastSeen = timeStamp
+	c.Sessions[id] = session
+
 	return true
 }
 
+// Returns the session ID that contains a given connection
+func (c *Channel) GetSessionIDByConn(conn net.Conn) (string, bool) {
+	if len(c.Sessions) == 0 {
+		output.PrintFrameworkDebug("No sessions exist")
+
+		return "", false
+	}
+
+	for id, session := range c.Sessions {
+		if *session.conn == conn {
+			return id, true
+		}
+	}
+
+	output.PrintFrameworkError("Conn does not exist in sessions")
+
+	return "", false
+}
+
+
 // RemoveSession removes a specific session ID and if a connection exists, closes it.
 func (c *Channel) RemoveSession(id string) bool {
 	if len(c.Sessions) == 0 {
@@ -77,7 +127,7 @@ func (c *Channel) RemoveSession(id string) bool {
 
 		return false
 	}
-	_, ok := c.Sessions[id]
+	session, ok := c.Sessions[id]
 	if !ok {
 		output.PrintFrameworkError("Session ID does not exist")
 
@@ -86,7 +136,8 @@ func (c *Channel) RemoveSession(id string) bool {
 	if c.Sessions[id].conn != nil {
 		(*c.Sessions[id].conn).Close()
 	}
-	delete(c.Sessions, id)
+	session.Active = false
+	c.Sessions[id] = session
 
 	return true
 }
@@ -98,11 +149,8 @@ func (c *Channel) RemoveSessions() bool {
 
 		return false
 	}
-	for k := range c.Sessions {
-		if c.Sessions[k].conn != nil {
-			(*c.Sessions[k].conn).Close()
-		}
-		delete(c.Sessions, k)
+	for id := range c.Sessions {
+		c.RemoveSession(id)
 	}
 
 	return true
diff --git a/c2/cli/basic.go b/c2/cli/basic.go
index 11d6c23..554bf8b 100644
--- a/c2/cli/basic.go
+++ b/c2/cli/basic.go
@@ -48,6 +48,13 @@ func backgroundResponse(ch *channel.Channel, wg *sync.WaitGroup, conn net.Conn,
 			// could have move data to write, but the user has already called exit
 			// below. I that that's tolerable for now.
 			responseCh <- string(responseBuffer[:bytesRead])
+			// Update "Last Seen"
+			ok := ch.UpdateLastSeenByConn(conn, time.Now())
+			if !ok {
+				output.PrintFrameworkError("Failed to update LastSeen value for connection")
+
+				return
+			}
 		}
 		time.Sleep(10 * time.Millisecond)
 	}