From f75b75790c6b56af1f3d93d64f3450c36320bffd Mon Sep 17 00:00:00 2001
From: sawka <mike.sawka@gmail.com>
Date: Tue, 20 Sep 2022 14:15:20 -0700
Subject: [PATCH] session stats/size

---
 pkg/sstore/dbops.go   | 28 +++++++++++++++++++
 pkg/sstore/fileops.go | 65 +++++++++++++++++++++++++++++++++++++++++++
 pkg/sstore/sstore.go  |  9 ++++++
 3 files changed, 102 insertions(+)

diff --git a/pkg/sstore/dbops.go b/pkg/sstore/dbops.go
index a999de565..c88dfd57c 100644
--- a/pkg/sstore/dbops.go
+++ b/pkg/sstore/dbops.go
@@ -925,3 +925,31 @@ func UpdateCmdTermOpts(ctx context.Context, sessionId string, cmdId string, term
 func DeleteSession(ctx context.Context, sessionId string) error {
 	return nil
 }
+
+func GetSessionStats(ctx context.Context, sessionId string) (*SessionStatsType, error) {
+	rtn := &SessionStatsType{SessionId: sessionId}
+	txErr := WithTx(ctx, func(tx *TxWrap) error {
+		query := `SELECT sessionid FROM session WHERE sessionid = ?`
+		if !tx.Exists(query, sessionId) {
+			return fmt.Errorf("not found")
+		}
+		query = `SELECT count(*) FROM screen WHERE sessionid = ?`
+		rtn.NumScreens = tx.GetInt(query, sessionId)
+		query = `SELECT count(*) FROM window WHERE sessionid = ?`
+		rtn.NumWindows = tx.GetInt(query, sessionId)
+		query = `SELECT count(*) FROM line WHERE sessionid = ?`
+		rtn.NumLines = tx.GetInt(query, sessionId)
+		query = `SELECT count(*) FROM cmd WHERE sessionid = ?`
+		rtn.NumCmds = tx.GetInt(query, sessionId)
+		return nil
+	})
+	if txErr != nil {
+		return nil, txErr
+	}
+	diskSize, err := SessionDiskSize(sessionId)
+	if err != nil {
+		return nil, err
+	}
+	rtn.DiskStats = diskSize
+	return rtn, nil
+}
diff --git a/pkg/sstore/fileops.go b/pkg/sstore/fileops.go
index 0ecf9aa02..1fdad2cd8 100644
--- a/pkg/sstore/fileops.go
+++ b/pkg/sstore/fileops.go
@@ -4,7 +4,11 @@ import (
 	"context"
 	"encoding/base64"
 	"fmt"
+	"os"
+	"path"
 
+	"github.com/google/uuid"
+	"github.com/scripthaus-dev/mshell/pkg/base"
 	"github.com/scripthaus-dev/mshell/pkg/cirfile"
 	"github.com/scripthaus-dev/sh2-server/pkg/scbase"
 )
@@ -62,3 +66,64 @@ func ReadFullPtyOutFile(ctx context.Context, sessionId string, cmdId string) (in
 	defer f.Close()
 	return f.ReadAll(ctx)
 }
+
+type SessionDiskSizeType struct {
+	NumFiles   int
+	TotalSize  int64
+	ErrorCount int
+}
+
+func directorySize(dirName string) (SessionDiskSizeType, error) {
+	var rtn SessionDiskSizeType
+	entries, err := os.ReadDir(dirName)
+	if err != nil {
+		return rtn, err
+	}
+	for _, entry := range entries {
+		if entry.IsDir() {
+			rtn.ErrorCount++
+			continue
+		}
+		finfo, err := entry.Info()
+		if err != nil {
+			rtn.ErrorCount++
+			continue
+		}
+		rtn.NumFiles++
+		rtn.TotalSize += finfo.Size()
+	}
+	return rtn, nil
+}
+
+func SessionDiskSize(sessionId string) (SessionDiskSizeType, error) {
+	sessionDir, err := base.EnsureSessionDir(sessionId)
+	if err != nil {
+		return SessionDiskSizeType{}, err
+	}
+	return directorySize(sessionDir)
+}
+
+func FullSessionDiskSize() (map[string]SessionDiskSizeType, error) {
+	sdir := base.GetSessionsDir()
+	entries, err := os.ReadDir(sdir)
+	if err != nil {
+		return nil, err
+	}
+	rtn := make(map[string]SessionDiskSizeType)
+	for _, entry := range entries {
+		if !entry.IsDir() {
+			continue
+		}
+		name := entry.Name()
+		_, err = uuid.Parse(name)
+		if err != nil {
+			continue
+		}
+		diskSize, err := directorySize(path.Join(sdir, name))
+		if err != nil {
+			continue
+		}
+		rtn[name] = diskSize
+	}
+	return rtn, nil
+}
diff --git a/pkg/sstore/sstore.go b/pkg/sstore/sstore.go
index 9efcbcae8..d1b073015 100644
--- a/pkg/sstore/sstore.go
+++ b/pkg/sstore/sstore.go
@@ -117,6 +117,15 @@ type SessionType struct {
 	Full   bool `json:"full,omitempty"`
 }
 
+type SessionStatsType struct {
+	SessionId  string              `json:"sessionid"`
+	NumScreens int                 `json:"numscreens"`
+	NumWindows int                 `json:"numwindows"`
+	NumLines   int                 `json:"numlines"`
+	NumCmds    int                 `json:"numcmds"`
+	DiskStats  SessionDiskSizeType `json:"diskstats"`
+}
+
 type WindowOptsType struct {
 }