mirror of
https://github.com/wavetermdev/waveterm.git
synced 2024-12-27 17:37:58 +01:00
369 lines
11 KiB
Go
369 lines
11 KiB
Go
package sstore
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"fmt"
|
|
"strings"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/scripthaus-dev/mshell/pkg/packet"
|
|
)
|
|
|
|
func NumSessions(ctx context.Context) (int, error) {
|
|
db, err := GetDB(ctx)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
query := "SELECT count(*) FROM session"
|
|
var count int
|
|
err = db.GetContext(ctx, &count, query)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
return count, nil
|
|
}
|
|
|
|
func GetAllRemotes(ctx context.Context) ([]*RemoteType, error) {
|
|
var rtn []*RemoteType
|
|
err := WithTx(ctx, func(tx *TxWrap) error {
|
|
query := `SELECT * FROM remote`
|
|
marr := tx.SelectMaps(query)
|
|
for _, m := range marr {
|
|
rtn = append(rtn, RemoteFromMap(m))
|
|
}
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return rtn, nil
|
|
}
|
|
|
|
func GetRemoteByName(ctx context.Context, remoteName string) (*RemoteType, error) {
|
|
var remote *RemoteType
|
|
err := WithTx(ctx, func(tx *TxWrap) error {
|
|
query := `SELECT * FROM remote WHERE remotename = ?`
|
|
m := tx.GetMap(query, remoteName)
|
|
remote = RemoteFromMap(m)
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return remote, nil
|
|
}
|
|
|
|
func GetRemoteById(ctx context.Context, remoteId string) (*RemoteType, error) {
|
|
var remote *RemoteType
|
|
err := WithTx(ctx, func(tx *TxWrap) error {
|
|
query := `SELECT * FROM remote WHERE remoteid = ?`
|
|
m := tx.GetMap(query, remoteId)
|
|
remote = RemoteFromMap(m)
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return remote, nil
|
|
}
|
|
|
|
func InsertRemote(ctx context.Context, remote *RemoteType) error {
|
|
if remote == nil {
|
|
return fmt.Errorf("cannot insert nil remote")
|
|
}
|
|
db, err := GetDB(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
query := `INSERT INTO remote ( remoteid, remotetype, remotename, autoconnect, initpk, sshopts, lastconnectts) VALUES
|
|
(:remoteid,:remotetype,:remotename,:autoconnect,:initpk,:sshopts,:lastconnectts)`
|
|
_, err = db.NamedExec(query, remote.ToMap())
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func GetAllSessions(ctx context.Context) ([]*SessionType, error) {
|
|
var rtn []*SessionType
|
|
err := WithTx(ctx, func(tx *TxWrap) error {
|
|
query := `SELECT * FROM session`
|
|
tx.SelectWrap(&rtn, query)
|
|
var windows []*WindowType
|
|
query = `SELECT * FROM window`
|
|
tx.SelectWrap(&windows, query)
|
|
winMap := make(map[string][]*WindowType)
|
|
for _, win := range windows {
|
|
winArr := winMap[win.SessionId]
|
|
winArr = append(winArr, win)
|
|
winMap[win.SessionId] = winArr
|
|
}
|
|
for _, session := range rtn {
|
|
session.Windows = winMap[session.SessionId]
|
|
}
|
|
var screens []*ScreenType
|
|
query = `SELECT * FROM screen ORDER BY screenidx`
|
|
tx.SelectWrap(&screens, query)
|
|
screenMap := make(map[string][]*ScreenType)
|
|
for _, screen := range screens {
|
|
screenArr := screenMap[screen.SessionId]
|
|
screenArr = append(screenArr, screen)
|
|
screenMap[screen.SessionId] = screenArr
|
|
}
|
|
for _, session := range rtn {
|
|
session.Screens = screenMap[session.SessionId]
|
|
}
|
|
var sws []*ScreenWindowType
|
|
query = `SELECT * FROM screen_window`
|
|
tx.SelectWrap(&sws, query)
|
|
screenIdMap := make(map[string]*ScreenType)
|
|
for _, screen := range screens {
|
|
screenIdMap[screen.SessionId+screen.ScreenId] = screen
|
|
}
|
|
for _, sw := range sws {
|
|
screen := screenIdMap[sw.SessionId+sw.ScreenId]
|
|
if screen == nil {
|
|
continue
|
|
}
|
|
screen.Windows = append(screen.Windows, sw)
|
|
}
|
|
return nil
|
|
})
|
|
return rtn, err
|
|
}
|
|
|
|
func GetWindowById(ctx context.Context, sessionId string, windowId string) (*WindowType, error) {
|
|
var rtnWindow *WindowType
|
|
err := WithTx(ctx, func(tx *TxWrap) error {
|
|
var window WindowType
|
|
query := `SELECT * FROM window WHERE sessionid = ? AND windowid = ?`
|
|
found := tx.GetWrap(&window, query, sessionId, windowId)
|
|
if !found {
|
|
return nil
|
|
}
|
|
rtnWindow = &window
|
|
query = `SELECT * FROM line WHERE sessionid = ? AND windowid = ?`
|
|
tx.SelectWrap(&window.Lines, query, sessionId, windowId)
|
|
query = `SELECT * FROM cmd WHERE cmdid IN (SELECT cmdid FROM line WHERE sessionid = ? AND windowid = ?)`
|
|
cmdMaps := tx.SelectMaps(query, sessionId, windowId)
|
|
for _, m := range cmdMaps {
|
|
window.Cmds = append(window.Cmds, CmdFromMap(m))
|
|
}
|
|
return nil
|
|
})
|
|
return rtnWindow, err
|
|
}
|
|
|
|
func GetSessionById(ctx context.Context, id string) (*SessionType, error) {
|
|
var rtnSession *SessionType
|
|
err := WithTx(ctx, func(tx *TxWrap) error {
|
|
var session SessionType
|
|
query := `SELECT * FROM session WHERE sessionid = ?`
|
|
found := tx.GetWrap(&session, query, id)
|
|
if !found {
|
|
return nil
|
|
}
|
|
rtnSession = &session
|
|
query = `SELECT * FROM window WHERE sessionid = ?`
|
|
tx.SelectWrap(&session.Windows, query, session.SessionId)
|
|
query = `SELECT * FROM remote_instance WHERE sessionid = ?`
|
|
tx.SelectWrap(&session.Remotes, query, session.SessionId)
|
|
query = `SELECT * FROM cmd WHERE sessionid = ?`
|
|
marr := tx.SelectMaps(query, session.SessionId)
|
|
for _, m := range marr {
|
|
session.Cmds = append(session.Cmds, CmdFromMap(m))
|
|
}
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return rtnSession, nil
|
|
}
|
|
|
|
func GetSessionByName(ctx context.Context, name string) (*SessionType, error) {
|
|
db, err := GetDB(ctx)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
var sessionId string
|
|
query := `SELECT sessionid FROM session WHERE name = ?`
|
|
err = db.GetContext(ctx, &sessionId, query, name)
|
|
if err != nil {
|
|
if err == sql.ErrNoRows {
|
|
return nil, nil
|
|
}
|
|
return nil, err
|
|
}
|
|
return GetSessionById(ctx, sessionId)
|
|
}
|
|
|
|
// also creates default window, returns sessionId
|
|
// if sessionName == "", it will be generated
|
|
func InsertSessionWithName(ctx context.Context, sessionName string) (string, error) {
|
|
newSessionId := uuid.New().String()
|
|
txErr := WithTx(ctx, func(tx *TxWrap) error {
|
|
names := tx.SelectStrings(`SELECT name FROM session`)
|
|
sessionName = fmtUniqueName(sessionName, "session-%d", len(names)+1, names)
|
|
maxSessionIdx := tx.GetInt(`SELECT COALESCE(max(sessionidx), 0) FROM session`)
|
|
query := `INSERT INTO session (sessionid, name, activescreenid, sessionidx, notifynum) VALUES (?, ?, '', ?, ?)`
|
|
tx.ExecWrap(query, newSessionId, sessionName, maxSessionIdx+1, 0)
|
|
screenId, err := InsertScreen(tx.Context(), newSessionId, "")
|
|
if err != nil {
|
|
return err
|
|
}
|
|
query = `UPDATE session SET activescreenid = ? WHERE sessionid = ?`
|
|
tx.ExecWrap(query, screenId, newSessionId)
|
|
return nil
|
|
})
|
|
return newSessionId, txErr
|
|
}
|
|
|
|
func containsStr(strs []string, testStr string) bool {
|
|
for _, s := range strs {
|
|
if s == testStr {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func fmtUniqueName(name string, defaultFmtStr string, startIdx int, strs []string) string {
|
|
var fmtStr string
|
|
if name != "" {
|
|
if !containsStr(strs, name) {
|
|
return name
|
|
}
|
|
fmtStr = name + "-%d"
|
|
startIdx = 2
|
|
} else {
|
|
fmtStr = defaultFmtStr
|
|
}
|
|
if strings.Index(fmtStr, "%d") == -1 {
|
|
panic("invalid fmtStr: " + fmtStr)
|
|
}
|
|
for {
|
|
testName := fmt.Sprintf(fmtStr, startIdx)
|
|
if containsStr(strs, testName) {
|
|
startIdx++
|
|
continue
|
|
}
|
|
return testName
|
|
}
|
|
}
|
|
|
|
func InsertScreen(ctx context.Context, sessionId string, screenName string) (string, error) {
|
|
var newScreenId string
|
|
txErr := WithTx(ctx, func(tx *TxWrap) error {
|
|
query := `SELECT sessionid FROM session WHERE sessionid = ?`
|
|
if !tx.Exists(query, sessionId) {
|
|
return fmt.Errorf("cannot create screen, no session found")
|
|
}
|
|
newWindowId := txCreateWindow(tx, sessionId)
|
|
maxScreenIdx := tx.GetInt(`SELECT COALESCE(max(screenidx), 0) FROM screen WHERE sessionid = ?`, sessionId)
|
|
screenNames := tx.SelectStrings(`SELECT name FROM screen WHERE sessionid = ?`, sessionId)
|
|
screenName = fmtUniqueName(screenName, "s%d", maxScreenIdx+1, screenNames)
|
|
newScreenId = uuid.New().String()
|
|
query = `INSERT INTO screen (sessionid, screenid, name, activewindowid, screenidx, screenopts) VALUES (?, ?, ?, ?, ?, ?)`
|
|
tx.ExecWrap(query, sessionId, newScreenId, screenName, newWindowId, maxScreenIdx+1, ScreenOptsType{})
|
|
layout := LayoutType{Type: LayoutFull}
|
|
query = `INSERT INTO screen_window (sessionid, screenid, windowid, name, layout) VALUES (?, ?, ?, ?, ?)`
|
|
tx.ExecWrap(query, sessionId, newScreenId, newWindowId, DefaultScreenWindowName, layout)
|
|
return nil
|
|
})
|
|
return newScreenId, txErr
|
|
}
|
|
|
|
func txCreateWindow(tx *TxWrap, sessionId string) string {
|
|
windowId := uuid.New().String()
|
|
query := `INSERT INTO window (sessionid, windowid, curremote, winopts) VALUES (?, ?, ?, ?)`
|
|
tx.ExecWrap(query, sessionId, windowId, LocalRemoteName, WindowOptsType{})
|
|
return windowId
|
|
}
|
|
|
|
func InsertLine(ctx context.Context, line *LineType, cmd *CmdType) error {
|
|
if line == nil {
|
|
return fmt.Errorf("line cannot be nil")
|
|
}
|
|
if line.LineId != 0 {
|
|
return fmt.Errorf("new line cannot have LineId set")
|
|
}
|
|
return WithTx(ctx, func(tx *TxWrap) error {
|
|
var windowId string
|
|
query := `SELECT windowid FROM window WHERE sessionid = ? AND windowid = ?`
|
|
hasWindow := tx.GetWrap(&windowId, query, line.SessionId, line.WindowId)
|
|
if !hasWindow {
|
|
return fmt.Errorf("window not found, cannot insert line[%s/%s]", line.SessionId, line.WindowId)
|
|
}
|
|
var maxLineId int64
|
|
query = `SELECT COALESCE(max(lineid), 0) FROM line WHERE sessionid = ? AND windowid = ?`
|
|
tx.GetWrap(&maxLineId, query, line.SessionId, line.WindowId)
|
|
line.LineId = maxLineId + 1
|
|
query = `INSERT INTO line ( sessionid, windowid, lineid, ts, userid, linetype, text, cmdid)
|
|
VALUES (:sessionid,:windowid,:lineid,:ts,:userid,:linetype,:text,:cmdid)`
|
|
tx.NamedExecWrap(query, line)
|
|
if cmd != nil {
|
|
cmdMap := cmd.ToMap()
|
|
query = `
|
|
INSERT INTO cmd ( sessionid, cmdid, remoteid, cmdstr, remotestate, termopts, status, startpk, donepk, runout, usedrows)
|
|
VALUES (:sessionid,:cmdid,:remoteid,:cmdstr,:remotestate,:termopts,:status,:startpk,:donepk,:runout,:usedrows)
|
|
`
|
|
tx.NamedExecWrap(query, cmdMap)
|
|
}
|
|
return nil
|
|
})
|
|
}
|
|
|
|
func GetCmdById(ctx context.Context, sessionId string, cmdId string) (*CmdType, error) {
|
|
var cmd *CmdType
|
|
err := WithTx(ctx, func(tx *TxWrap) error {
|
|
query := `SELECT * FROM cmd WHERE sessionid = ? AND cmdid = ?`
|
|
m := tx.GetMap(query, sessionId, cmdId)
|
|
cmd = CmdFromMap(m)
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return cmd, nil
|
|
}
|
|
|
|
func UpdateCmdDonePk(ctx context.Context, donePk *packet.CmdDonePacketType) error {
|
|
if donePk == nil || donePk.CK.IsEmpty() {
|
|
return fmt.Errorf("invalid cmddone packet (no ck)")
|
|
}
|
|
return WithTx(ctx, func(tx *TxWrap) error {
|
|
query := `UPDATE cmd SET status = ?, donepk = ? WHERE sessionid = ? AND cmdid = ?`
|
|
tx.ExecWrap(query, CmdStatusDone, quickJson(donePk), donePk.CK.GetSessionId(), donePk.CK.GetCmdId())
|
|
return nil
|
|
})
|
|
}
|
|
|
|
func AppendCmdErrorPk(ctx context.Context, errPk *packet.CmdErrorPacketType) error {
|
|
if errPk == nil || errPk.CK.IsEmpty() {
|
|
return fmt.Errorf("invalid cmderror packet (no ck)")
|
|
}
|
|
return WithTx(ctx, func(tx *TxWrap) error {
|
|
query := `UPDATE cmd SET runout = json_insert(runout, '$[#]', ?) WHERE sessionid = ? AND cmdid = ?`
|
|
tx.ExecWrap(query, quickJson(errPk), errPk.CK.GetSessionId(), errPk.CK.GetCmdId())
|
|
return nil
|
|
})
|
|
}
|
|
|
|
func HangupAllRunningCmds(ctx context.Context) error {
|
|
return WithTx(ctx, func(tx *TxWrap) error {
|
|
query := `UPDATE cmd SET status = ? WHERE status = ?`
|
|
tx.ExecWrap(query, CmdStatusHangup, CmdStatusRunning)
|
|
return nil
|
|
})
|
|
}
|
|
|
|
func HangupRunningCmdsByRemoteId(ctx context.Context, remoteId string) error {
|
|
return WithTx(ctx, func(tx *TxWrap) error {
|
|
query := `UPDATE cmd SET status = ? WHERE status = ? AND remoteid = ?`
|
|
tx.ExecWrap(query, CmdStatusHangup, CmdStatusRunning, remoteId)
|
|
return nil
|
|
})
|
|
}
|