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) }