converge/pkg/agent/session.go

package agent

import (
	"fmt"
	"github.com/gliderlabs/ssh"
	"io"
	"log"
	"os"
	"strconv"
	"strings"
	"time"

	_ "embed"
)

// global configuration

type AgentState struct {
	startTime time.Time

	// Advance warning time to notify the user of something important happening
	advanceWarningTime time.Duration

	// session expiry time
	agentExpriryTime time.Duration

	// ticker
	tickerInterval time.Duration
	ticker         *time.Ticker

	// map of unique session id to a session
	sessions map[int]*AgentSession

	agentUsed bool
}

type AgentSession struct {
	startTime time.Time

	// For sending messages to the user
	sshSession ssh.Session
}

var state AgentState

const holdFilename = ".hold"

//go:embed help.txt
var helpMessage string

func ConfigureAgent(advanceWarningTime, agentExpiryTime, tickerInterval time.Duration) {
	if fileExists(holdFilename) {
		log.Printf("Removing hold file '%s'", holdFilename)
		err := os.Remove(holdFilename)
		if err != nil {
			log.Printf("Could not remove hold file: '%s'", holdFilename)
		}
	}
	state = AgentState{
		startTime:          time.Now(),
		advanceWarningTime: advanceWarningTime,
		agentExpriryTime:   agentExpiryTime,
		tickerInterval:     tickerInterval,
		ticker:             time.NewTicker(tickerInterval),
		sessions:           make(map[int]*AgentSession),
		agentUsed:          false,
	}

	go func() {
		for {
			<-state.ticker.C
			check()
		}
	}()
}

func Login(sessionId int, sshSession ssh.Session) {
	log.Println("New login")
	hostname, _ := os.Hostname()

	holdFileStats, ok := fileExistsWithStats(holdFilename)
	if ok {
		if holdFileStats.ModTime().After(time.Now()) {
			// modification time in the future, leaving intact
			log.Println("Hold file has modification time in the future, leaving intact")
		} else {
			log.Printf("Touching hold file '%s'", holdFilename)
			err := os.Chtimes(holdFilename, time.Now(), time.Now())
			if err != nil {
				log.Printf("Could not touch hold file: '%s'", holdFilename)
			}
		}
	}

	PrintMessage(sshSession, fmt.Sprintf("You are now on %s\n", hostname))
	PrintHelpMessage(sshSession)

	agentSession := AgentSession{
		startTime:  time.Now(),
		sshSession: sshSession,
	}
	state.sessions[sessionId] = &agentSession
	state.agentUsed = true
	LogStatus()
}

func PrintHelpMessage(sshSession ssh.Session) {
	PrintMessage(sshSession, fmt.Sprintf(helpMessage,
		state.expiryTime(holdFilename).Format(time.DateTime),
		state.agentExpriryTime))
}

func LogOut(sessionId int) {
	log.Println("User logged out")
	delete(state.sessions, sessionId)
	LogStatus()
	check()
}

func PrintMessage(sshSession ssh.Session, message string) {
	io.WriteString(sshSession.Stderr(), "\n\r###\n\r")
	for _, line := range strings.Split(message, "\n") {
		io.WriteString(sshSession.Stderr(), "### "+line+"\n\r")
	}
	io.WriteString(sshSession.Stderr(), "\n\r")
}

func LogStatus() {
	fmt := "%-20s %-20s %-20s"
	log.Println()
	log.Printf(fmt, "UID", "START_TIME", "TYPE")
	for uid, session := range state.sessions {
		sessionType := session.sshSession.Subsystem()
		if sessionType == "" {
			sessionType = "ssh"
		}
		log.Printf(fmt, strconv.Itoa(uid),
			session.startTime.Format(time.DateTime),
			sessionType)
	}
	log.Println()
}

func fileExistsWithStats(filename string) (os.FileInfo, bool) {
	stats, err := os.Stat(filename)
	return stats, !os.IsNotExist(err)
}

func fileExists(filename string) bool {
	_, err := os.Stat(filename)
	return !os.IsNotExist(err)
}

func (state *AgentState) expiryTime(filename string) time.Time {
	stats, err := os.Stat(filename)
	if err != nil {
		return state.startTime.Add(state.agentExpriryTime)
	}
	return stats.ModTime().Add(state.agentExpriryTime)
}

// Behavior to implement
//  1. there is a global timeout for all agent sessions together: state.agentExpirtyTime
//  2. The expiry time is relative to the modification time of the .hold file in the
//     agent directory or, if that file does not exist, the start time of the agent.
//  3. if we are close to the expiry time then we message users with instruction on
//     how to prevent the timeout
//  4. If the last user logs out, the aagent will exit immediately if no .hold file is
//     present. Otherwise it will exit after the epxiry time. This allows users to
//     reconnect later.
func check() {
	now := time.Now()

	expiryTime := state.expiryTime(".hold")

	if now.After(expiryTime) {
		messageUsers("Expiry time was reached logging out")
		time.Sleep(5 * time.Second)
		log.Println("Agent exiting")
		os.Exit(0)
	}

	if expiryTime.Sub(now) < state.advanceWarningTime {
		messageUsers(
			fmt.Sprintf("Session will expire at %s", expiryTime.Format(time.DateTime)))
		for _, session := range state.sessions {
			PrintHelpMessage(session.sshSession)
		}
	}

	if state.agentUsed && !fileExists(holdFilename) && len(state.sessions) == 0 {
		log.Printf("All clients disconnected and no '%s' file found, exiting", holdFilename)
		os.Exit(0)
	}
}

func messageUsers(message string) {
	log.Printf("=== Notification to users: %s", message)
	for _, session := range state.sessions {
		PrintMessage(session.sshSession, message)
	}
}