From 134c72d8d01647c0b1e396fac9ff75c46dc472bc Mon Sep 17 00:00:00 2001
From: Erik Brakkee <erik@brakkee.org>
Date: Sun, 21 Jul 2024 12:34:34 +0200
Subject: [PATCH] basic session management is now implemented.

---
 cmd/agent/agent.go             | 13 +++++
 pkg/agent/session.go           | 87 ++++++++++++++++++++++++++++++++++
 pkg/{sshutils => agent}/ssh.go |  2 +-
 pkg/converge/admin.go          |  2 +-
 4 files changed, 102 insertions(+), 2 deletions(-)
 create mode 100644 pkg/agent/session.go
 rename pkg/{sshutils => agent}/ssh.go (97%)

diff --git a/cmd/agent/agent.go b/cmd/agent/agent.go
index ddc70c3..6d663cc 100755
--- a/cmd/agent/agent.go
+++ b/cmd/agent/agent.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"bufio"
+	"cidebug/pkg/agent"
 	"cidebug/pkg/iowrappers"
 	"cidebug/pkg/websocketutil"
 	"fmt"
@@ -12,6 +13,7 @@ import (
 	"os"
 	"os/exec"
 	"syscall"
+	"time"
 	"unsafe"
 
 	"github.com/creack/pty"
@@ -73,11 +75,16 @@ func sshServer(hostKeyFile string) *ssh.Server {
 					setWinsize(f, win.Width, win.Height)
 				}
 			}()
+			uid := int(time.Now().UnixMilli())
+			agent.Login(uid, s)
+
 			go func() {
 				io.Copy(f, s) // stdin
 			}()
 			io.Copy(s, f) // stdout
 			cmd.Wait()
+			log.Println("User logged out")
+			agent.LogOut(uid)
 		} else {
 			io.WriteString(s, "No PTY requested.\n")
 			s.Exit(1)
@@ -145,6 +152,12 @@ func (f ReaderFunc) Read(p []byte) (n int, err error) {
 
 func main() {
 	wsURL := os.Args[1]
+
+	advanceWarningTime := 1 * time.Minute
+	sessionExpiryTime := 5 * time.Minute
+	tickerInterval := 10 * time.Second
+	agent.ConfigureAgent(advanceWarningTime, sessionExpiryTime, tickerInterval)
+
 	conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
 	if err != nil {
 		log.Println("WebSocket connection error:", err)
diff --git a/pkg/agent/session.go b/pkg/agent/session.go
new file mode 100644
index 0000000..4511269
--- /dev/null
+++ b/pkg/agent/session.go
@@ -0,0 +1,87 @@
+package agent
+
+import (
+	"github.com/gliderlabs/ssh"
+	"io"
+	"log"
+	"strconv"
+	"time"
+)
+
+// global configuration
+
+type AgentState struct {
+
+	// Advance warning time to notify the user of something important happening
+	advanceWarningTime time.Duration
+
+	// session expiry time
+	sessionExpiryTime time.Duration
+
+	// ticker
+	tickerInterval time.Duration
+	ticker         *time.Ticker
+
+	// map of unique session id to a session
+	sessions map[int]*AgentSession
+}
+
+type AgentSession struct {
+	startTime time.Time
+
+	// For sending messages to the user
+	sshSession ssh.Session
+}
+
+var state AgentState
+
+func ConfigureAgent(advanceWarningTime, sessionExpiryTime, tickerInterval time.Duration) {
+	state = AgentState{
+		advanceWarningTime: advanceWarningTime,
+		sessionExpiryTime:  sessionExpiryTime,
+		tickerInterval:     tickerInterval,
+		ticker:             time.NewTicker(tickerInterval),
+		sessions:           make(map[int]*AgentSession),
+	}
+
+	go func() {
+		for {
+			<-state.ticker.C
+			check()
+		}
+	}()
+}
+
+func Login(sessionId int, sshSession ssh.Session) {
+	log.Println("New login")
+	agentSession := AgentSession{
+		startTime:  time.Now(),
+		sshSession: sshSession,
+	}
+	state.sessions[sessionId] = &agentSession
+	LogStatus()
+}
+
+func LogOut(sessionId int) {
+	log.Println("User logged out")
+	delete(state.sessions, sessionId)
+	LogStatus()
+	check()
+}
+
+func LogStatus() {
+	fmt := "%-20s %-20s"
+	log.Println()
+	log.Printf(fmt, "UID", "START_TIME")
+	for uid, session := range state.sessions {
+		log.Printf(fmt, strconv.Itoa(uid), session.startTime.Format("2006-01-02 15:04:05"))
+	}
+	log.Println()
+}
+
+func check() {
+	log.Println("Timer is firing!")
+	for _, session := range state.sessions {
+		io.WriteString(session.sshSession.Stderr(), "\n\nThe clock is ticking for you!\n\n")
+	}
+}
diff --git a/pkg/sshutils/ssh.go b/pkg/agent/ssh.go
similarity index 97%
rename from pkg/sshutils/ssh.go
rename to pkg/agent/ssh.go
index 3111c82..53f655f 100644
--- a/pkg/sshutils/ssh.go
+++ b/pkg/agent/ssh.go
@@ -1,4 +1,4 @@
-package sshutils
+package agent
 
 import (
 	"crypto/rand"
diff --git a/pkg/converge/admin.go b/pkg/converge/admin.go
index b53d9cf..5372d8d 100644
--- a/pkg/converge/admin.go
+++ b/pkg/converge/admin.go
@@ -169,7 +169,7 @@ func (admin *Admin) Register(publicId string, conn io.ReadWriteCloser) error {
 	}()
 	log.Printf("Agent registered: '%s'\n", publicId)
 	for !agent.clientSession.IsClosed() {
-		time.Sleep(1 * time.Second)
+		time.Sleep(250 * time.Millisecond)
 	}
 	return nil
 }