From ab5deafdb6fb8d8305b2a65c86000c9a84ba4f44 Mon Sep 17 00:00:00 2001 From: sawka Date: Tue, 2 May 2023 12:43:54 -0700 Subject: [PATCH] updates to remote table, working on openai remote type --- .gitignore | 3 ++ db/migrations/000018_modremote.down.sql | 14 +++++ db/migrations/000018_modremote.up.sql | 11 ++++ go.mod | 2 + go.sum | 2 + pkg/cmdrunner/cmdrunner.go | 11 ++-- pkg/remote/remote.go | 25 +++++---- pkg/sstore/dbops.go | 19 ++----- pkg/sstore/migrate.go | 2 +- pkg/sstore/sstore.go | 69 +++++++++++-------------- 10 files changed, 88 insertions(+), 70 deletions(-) create mode 100644 db/migrations/000018_modremote.down.sql create mode 100644 db/migrations/000018_modremote.up.sql diff --git a/.gitignore b/.gitignore index a89c46b43..7c9c1bf40 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,6 @@ *~ bin/ *.out +.idea/ +test/ +temp.sql diff --git a/db/migrations/000018_modremote.down.sql b/db/migrations/000018_modremote.down.sql new file mode 100644 index 000000000..d3d735349 --- /dev/null +++ b/db/migrations/000018_modremote.down.sql @@ -0,0 +1,14 @@ +ALTER TABLE remote ADD COLUMN remotesudo; + +UPDATE remote +SET remotesudo = 1 +WHERE json_extract(sshopts, '$.issudo') +; + +UPDATE remote +SET sshopts = json_remove(sshopts, '$.issudo') +; + +ALTER TABLE remote ADD COLUMN physicalid varchar(36) NOT NULL DEFAULT ''; + +ALTER TABLE remote DROP COLUMN openaiopts; diff --git a/db/migrations/000018_modremote.up.sql b/db/migrations/000018_modremote.up.sql new file mode 100644 index 000000000..4cad763fe --- /dev/null +++ b/db/migrations/000018_modremote.up.sql @@ -0,0 +1,11 @@ +UPDATE remote +SET sshopts = json_set(sshopts, '$.issudo', json('true')) +WHERE remotesudo +; + +ALTER TABLE remote DROP COLUMN remotesudo; + +ALTER TABLE remote DROP COLUMN physicalid; + +ALTER TABLE remote ADD COLUMN openaiopts json NOT NULL DEFAULT '{}'; + diff --git a/go.mod b/go.mod index 7803c8bf8..5c2dade3b 100644 --- a/go.mod +++ b/go.mod @@ -23,6 +23,8 @@ require ( require ( github.com/hashicorp/errwrap v1.1.0 // indirect github.com/hashicorp/go-multierror v1.1.1 // indirect + github.com/sashabaranov/go-openai v1.9.0 // indirect go.uber.org/atomic v1.7.0 // indirect ) + diff --git a/go.sum b/go.sum index 71eb2da2a..e7586676f 100644 --- a/go.sum +++ b/go.sum @@ -965,6 +965,8 @@ github.com/ruudk/golang-pdf417 v0.0.0-20181029194003-1af4ab5afa58/go.mod h1:6lfF github.com/ryanuber/columnize v0.0.0-20160712163229-9b3edd62028f/go.mod h1:sm1tb6uqfes/u+d4ooFouqFdy9/2g9QGwK3SQygK0Ts= github.com/safchain/ethtool v0.0.0-20190326074333-42ed695e3de8/go.mod h1:Z0q5wiBQGYcxhMZ6gUqHn6pYNLypFAvaL3UvgZLR0U4= github.com/safchain/ethtool v0.0.0-20210803160452-9aa261dae9b1/go.mod h1:Z0q5wiBQGYcxhMZ6gUqHn6pYNLypFAvaL3UvgZLR0U4= +github.com/sashabaranov/go-openai v1.9.0 h1:NoiO++IISxxJ1pRc0n7uZvMGMake0G+FJ1XPwXtprsA= +github.com/sashabaranov/go-openai v1.9.0/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg= github.com/satori/go.uuid v1.2.0/go.mod h1:dA0hQrYB0VpLJoorglMZABFdXlWrHn1NEOzdhQKdks0= github.com/sawka/txwrap v0.1.0 h1:uWGplmEJUEd9WGYZy9fU+hoC2Z6Yal4NMH5DbKsUTdo= github.com/sawka/txwrap v0.1.0/go.mod h1:T3nlw2gVpuolo6/XEetvBbk1oMXnY978YmBFy1UyHvw= diff --git a/pkg/cmdrunner/cmdrunner.go b/pkg/cmdrunner/cmdrunner.go index 2ec52098b..5b159e05b 100644 --- a/pkg/cmdrunner/cmdrunner.go +++ b/pkg/cmdrunner/cmdrunner.go @@ -544,14 +544,16 @@ func EvalCommand(ctx context.Context, pk *scpacket.FeCommandPacketType) (sstore. if len(pk.Args[0]) > MaxCommandLen { return nil, fmt.Errorf("command length too long len:%d, max:%d", len(pk.Args[0]), MaxCommandLen) } - if pk.Interactive { + evalDepth := getEvalDepth(ctx) + if pk.Interactive && evalDepth == 0 { err := sstore.UpdateCurrentActivity(ctx, sstore.ActivityUpdate{NumCommands: 1}) if err != nil { log.Printf("[error] incrementing activity numcommands: %v\n", err) // fall through (non-fatal error) } + log.Printf("inc numcommands\n") } - if getEvalDepth(ctx) > MaxEvalDepth { + if evalDepth > MaxEvalDepth { return nil, fmt.Errorf("alias/history expansion max-depth exceeded") } var historyContext historyContextType @@ -893,7 +895,6 @@ func makeRemoteEditErrorReturn_edit(ids resolvedIds, visual bool, err error) (ss type RemoteEditArgs struct { CanonicalName string SSHOpts *sstore.SSHOpts - Sudo bool ConnectMode string Alias string AutoInstall bool @@ -942,6 +943,7 @@ func parseRemoteEditArgs(isNew bool, pk *scpacket.FeCommandPacketType, isLocal b Local: false, SSHHost: remoteHost, SSHUser: remoteUser, + IsSudo: isSudo, } portVal, err := resolvePosInt(pk.Kwargs["port"], 0) if err != nil { @@ -1036,7 +1038,6 @@ func parseRemoteEditArgs(isNew bool, pk *scpacket.FeCommandPacketType, isLocal b return &RemoteEditArgs{ SSHOpts: sshOpts, - Sudo: isSudo, ConnectMode: connectMode, Alias: alias, AutoInstall: autoInstall, @@ -1060,11 +1061,9 @@ func RemoteNewCommand(ctx context.Context, pk *scpacket.FeCommandPacketType) (ss } r := &sstore.RemoteType{ RemoteId: scbase.GenPromptUUID(), - PhysicalId: "", RemoteType: sstore.RemoteTypeSsh, RemoteAlias: editArgs.Alias, RemoteCanonicalName: editArgs.CanonicalName, - RemoteSudo: editArgs.Sudo, RemoteUser: editArgs.SSHOpts.SSHUser, RemoteHost: editArgs.SSHOpts.SSHHost, ConnectMode: editArgs.ConnectMode, diff --git a/pkg/remote/remote.go b/pkg/remote/remote.go index 592c298b8..d0948be91 100644 --- a/pkg/remote/remote.go +++ b/pkg/remote/remote.go @@ -128,7 +128,6 @@ type RunCmdType struct { type RemoteRuntimeState struct { RemoteType string `json:"remotetype"` RemoteId string `json:"remoteid"` - PhysicalId string `json:"physicalremoteid"` RemoteAlias string `json:"remotealias,omitempty"` RemoteCanonicalName string `json:"remotecanonicalname"` RemoteVars map[string]string `json:"remotevars"` @@ -150,12 +149,22 @@ type RemoteRuntimeState struct { WaitingForPassword bool `json:"waitingforpassword,omitempty"` Local bool `json:"local,omitempty"` RemoteOpts *sstore.RemoteOptsType `json:"remoteopts,omitempty"` + CanComplete bool `json:"cancomplete,omitempty"` } func (state RemoteRuntimeState) IsConnected() bool { return state.Status == StatusConnected } +func CanComplete(remoteType string) bool { + switch remoteType { + case sstore.RemoteTypeSsh: + return true + default: + return false + } +} + func (msh *MShellProc) GetStatus() string { msh.Lock.Lock() defer msh.Lock.Unlock() @@ -234,7 +243,7 @@ func LoadRemotes(ctx context.Context) error { go msh.Launch(false) } if remote.Local { - if remote.RemoteSudo { + if remote.IsSudo() { numSudoLocal++ } else { numLocal++ @@ -498,7 +507,7 @@ func (msh *MShellProc) IsLocal() bool { func (msh *MShellProc) IsSudo() bool { msh.Lock.Lock() defer msh.Lock.Unlock() - return msh.Remote.RemoteSudo + return msh.Remote.IsSudo() } func (msh *MShellProc) tryAutoInstall() { @@ -519,7 +528,6 @@ func (msh *MShellProc) GetRemoteRuntimeState() RemoteRuntimeState { RemoteId: msh.Remote.RemoteId, RemoteAlias: msh.Remote.RemoteAlias, RemoteCanonicalName: msh.Remote.RemoteCanonicalName, - PhysicalId: msh.Remote.PhysicalId, Status: msh.Status, ConnectMode: msh.Remote.ConnectMode, AutoInstall: msh.Remote.AutoInstall, @@ -564,11 +572,10 @@ func (msh *MShellProc) GetRemoteRuntimeState() RemoteRuntimeState { vars["shorthost"] = makeShortHost(msh.Remote.RemoteHost) vars["alias"] = msh.Remote.RemoteAlias vars["cname"] = msh.Remote.RemoteCanonicalName - vars["physicalid"] = msh.Remote.PhysicalId vars["remoteid"] = msh.Remote.RemoteId vars["status"] = msh.Status vars["type"] = msh.Remote.RemoteType - if msh.Remote.RemoteSudo { + if msh.Remote.IsSudo() { vars["sudo"] = "1" } if msh.Remote.Local { @@ -603,9 +610,9 @@ func (msh *MShellProc) GetRemoteRuntimeState() RemoteRuntimeState { state.DefaultFeState = sstore.FeStateFromShellState(curState) vars["cwd"] = curState.Cwd } - if msh.Remote.Local && msh.Remote.RemoteSudo { + if msh.Remote.Local && msh.Remote.IsSudo() { vars["bestuser"] = "sudo" - } else if msh.Remote.RemoteSudo { + } else if msh.Remote.IsSudo() { vars["bestuser"] = "sudo@" + vars["bestuser"] } if msh.Remote.Local { @@ -1185,7 +1192,7 @@ func (msh *MShellProc) Launch(interactive bool) { var cmdStr string if sshOpts.SSHHost == "" && remoteCopy.Local { var err error - cmdStr, err = MakeLocalMShellCommandStr(remoteCopy.RemoteSudo) + cmdStr, err = MakeLocalMShellCommandStr(remoteCopy.IsSudo()) if err != nil { msh.WriteToPtyBuffer("*error, cannot find local mshell binary: %v\n", err) return diff --git a/pkg/sstore/dbops.go b/pkg/sstore/dbops.go index 71debe688..38977c5c5 100644 --- a/pkg/sstore/dbops.go +++ b/pkg/sstore/dbops.go @@ -165,19 +165,6 @@ func GetRemoteByCanonicalName(ctx context.Context, cname string) (*RemoteType, e return remote, nil } -func GetRemoteByPhysicalId(ctx context.Context, physicalId string) (*RemoteType, error) { - var remote *RemoteType - err := WithTx(ctx, func(tx *TxWrap) error { - query := `SELECT * FROM remote WHERE physicalid = ?` - remote = dbutil.GetMapGen[*RemoteType](tx, query, physicalId) - return nil - }) - if err != nil { - return nil, err - } - return remote, nil -} - func UpsertRemote(ctx context.Context, r *RemoteType) error { if r == nil { return fmt.Errorf("cannot insert nil remote") @@ -208,8 +195,8 @@ func UpsertRemote(ctx context.Context, r *RemoteType) error { maxRemoteIdx := tx.GetInt(query) r.RemoteIdx = int64(maxRemoteIdx + 1) query = `INSERT INTO remote - ( remoteid, physicalid, remotetype, remotealias, remotecanonicalname, remotesudo, remoteuser, remotehost, connectmode, autoinstall, sshopts, remoteopts, lastconnectts, archived, remoteidx, local, statevars) VALUES - (:remoteid,:physicalid,:remotetype,:remotealias,:remotecanonicalname,:remotesudo,:remoteuser,:remotehost,:connectmode,:autoinstall,:sshopts,:remoteopts,:lastconnectts,:archived,:remoteidx,:local,:statevars)` + ( remoteid, remotetype, remotealias, remotecanonicalname, remoteuser, remotehost, connectmode, autoinstall, sshopts, remoteopts, lastconnectts, archived, remoteidx, local, statevars, openaiopts) VALUES + (:remoteid,:remotetype,:remotealias,:remotecanonicalname,:remoteuser,:remotehost,:connectmode,:autoinstall,:sshopts,:remoteopts,:lastconnectts,:archived,:remoteidx,:local,:statevars,:openaiopts)` tx.NamedExec(query, r.ToMap()) return nil }) @@ -816,7 +803,7 @@ func InsertLine(ctx context.Context, line *LineType, cmd *CmdType) error { if line.LineNum != 0 { return fmt.Errorf("line should not hage linenum set") } - if cmd.ScreenId == "" { + if cmd != nil && cmd.ScreenId == "" { return fmt.Errorf("cmd should have screenid set") } return WithTx(ctx, func(tx *TxWrap) error { diff --git a/pkg/sstore/migrate.go b/pkg/sstore/migrate.go index fb1bfb7f7..f1a7b665c 100644 --- a/pkg/sstore/migrate.go +++ b/pkg/sstore/migrate.go @@ -17,7 +17,7 @@ import ( "github.com/golang-migrate/migrate/v4" ) -const MaxMigration = 17 +const MaxMigration = 18 const MigratePrimaryScreenVersion = 9 func MakeMigrate() (*migrate.Migrate, error) { diff --git a/pkg/sstore/sstore.go b/pkg/sstore/sstore.go index c969093d0..026588958 100644 --- a/pkg/sstore/sstore.go +++ b/pkg/sstore/sstore.go @@ -20,7 +20,6 @@ import ( "github.com/google/uuid" "github.com/jmoiron/sqlx" "github.com/sawka/txwrap" - "github.com/scripthaus-dev/mshell/pkg/base" "github.com/scripthaus-dev/mshell/pkg/packet" "github.com/scripthaus-dev/mshell/pkg/shexec" "github.com/scripthaus-dev/sh2-server/pkg/dbutil" @@ -76,7 +75,8 @@ const ( ) const ( - RemoteTypeSsh = "ssh" + RemoteTypeSsh = "ssh" + RemoteTypeOpenAI = "openai" ) const ( @@ -802,6 +802,7 @@ type ResolveItem struct { type SSHOpts struct { Local bool `json:"local,omitempty"` + IsSudo bool `json:"issudo,omitempty"` SSHHost string `json:"sshhost"` SSHUser string `json:"sshuser"` SSHOptsStr string `json:"sshopts,omitempty"` @@ -827,32 +828,34 @@ type RemoteOptsType struct { Color string `json:"color"` } -func (opts *RemoteOptsType) Scan(val interface{}) error { - return quickScanJson(opts, val) -} - -func (opts RemoteOptsType) Value() (driver.Value, error) { - return quickValueJson(opts) +type OpenAIOptsType struct { } type RemoteType struct { - RemoteId string `json:"remoteid"` - PhysicalId string `json:"physicalid"` - RemoteType string `json:"remotetype"` - RemoteAlias string `json:"remotealias"` - RemoteCanonicalName string `json:"remotecanonicalname"` - RemoteSudo bool `json:"remotesudo"` - RemoteUser string `json:"remoteuser"` - RemoteHost string `json:"remotehost"` - ConnectMode string `json:"connectmode"` - AutoInstall bool `json:"autoinstall"` - SSHOpts *SSHOpts `json:"sshopts"` - RemoteOpts *RemoteOptsType `json:"remoteopts"` - LastConnectTs int64 `json:"lastconnectts"` - Archived bool `json:"archived"` - RemoteIdx int64 `json:"remoteidx"` - Local bool `json:"local"` - StateVars map[string]string `json:"statevars"` + RemoteId string `json:"remoteid"` + RemoteType string `json:"remotetype"` + RemoteAlias string `json:"remotealias"` + RemoteCanonicalName string `json:"remotecanonicalname"` + RemoteOpts *RemoteOptsType `json:"remoteopts"` + LastConnectTs int64 `json:"lastconnectts"` + RemoteIdx int64 `json:"remoteidx"` + Archived bool `json:"archived"` + + // SSH fields + Local bool `json:"local"` + RemoteUser string `json:"remoteuser"` + RemoteHost string `json:"remotehost"` + ConnectMode string `json:"connectmode"` + AutoInstall bool `json:"autoinstall"` + SSHOpts *SSHOpts `json:"sshopts"` + StateVars map[string]string `json:"statevars"` + + // OpenAI fields + OpenAIOpts *OpenAIOptsType `json:"openaiopts,omitempty"` +} + +func (r *RemoteType) IsSudo() bool { + return r.SSHOpts != nil && r.SSHOpts.IsSudo } func (r *RemoteType) GetName() string { @@ -896,11 +899,9 @@ type CmdType struct { func (r *RemoteType) ToMap() map[string]interface{} { rtn := make(map[string]interface{}) rtn["remoteid"] = r.RemoteId - rtn["physicalid"] = r.PhysicalId rtn["remotetype"] = r.RemoteType rtn["remotealias"] = r.RemoteAlias rtn["remotecanonicalname"] = r.RemoteCanonicalName - rtn["remotesudo"] = r.RemoteSudo rtn["remoteuser"] = r.RemoteUser rtn["remotehost"] = r.RemoteHost rtn["connectmode"] = r.ConnectMode @@ -912,16 +913,15 @@ func (r *RemoteType) ToMap() map[string]interface{} { rtn["remoteidx"] = r.RemoteIdx rtn["local"] = r.Local rtn["statevars"] = quickJson(r.StateVars) + rtn["openaiopts"] = quickJson(r.OpenAIOpts) return rtn } func (r *RemoteType) FromMap(m map[string]interface{}) bool { quickSetStr(&r.RemoteId, m, "remoteid") - quickSetStr(&r.PhysicalId, m, "physicalid") quickSetStr(&r.RemoteType, m, "remotetype") quickSetStr(&r.RemoteAlias, m, "remotealias") quickSetStr(&r.RemoteCanonicalName, m, "remotecanonicalname") - quickSetBool(&r.RemoteSudo, m, "remotesudo") quickSetStr(&r.RemoteUser, m, "remoteuser") quickSetStr(&r.RemoteHost, m, "remotehost") quickSetStr(&r.ConnectMode, m, "connectmode") @@ -933,6 +933,7 @@ func (r *RemoteType) FromMap(m map[string]interface{}) bool { quickSetInt64(&r.RemoteIdx, m, "remoteidx") quickSetBool(&r.Local, m, "local") quickSetJson(&r.StateVars, m, "statevars") + quickSetJson(&r.OpenAIOpts, m, "openaiopts") return true } @@ -1029,10 +1030,6 @@ func AddCmdLine(ctx context.Context, screenId string, userId string, cmd *CmdTyp } func EnsureLocalRemote(ctx context.Context) error { - physicalId, err := base.GetRemoteId() - if err != nil { - return fmt.Errorf("getting local physical remoteid: %w", err) - } remote, err := GetLocalRemote(ctx) if err != nil { return fmt.Errorf("getting local remote from db: %w", err) @@ -1051,11 +1048,9 @@ func EnsureLocalRemote(ctx context.Context) error { // create the local remote localRemote := &RemoteType{ RemoteId: scbase.GenPromptUUID(), - PhysicalId: physicalId, RemoteType: RemoteTypeSsh, RemoteAlias: LocalRemoteAlias, RemoteCanonicalName: fmt.Sprintf("%s@%s", user.Username, hostName), - RemoteSudo: false, RemoteUser: user.Username, RemoteHost: hostName, ConnectMode: ConnectModeStartup, @@ -1070,16 +1065,14 @@ func EnsureLocalRemote(ctx context.Context) error { log.Printf("[db] added local remote '%s', id=%s\n", localRemote.RemoteCanonicalName, localRemote.RemoteId) sudoRemote := &RemoteType{ RemoteId: scbase.GenPromptUUID(), - PhysicalId: "", RemoteType: RemoteTypeSsh, RemoteAlias: "sudo", RemoteCanonicalName: fmt.Sprintf("sudo@%s@%s", user.Username, hostName), - RemoteSudo: true, RemoteUser: "root", RemoteHost: hostName, ConnectMode: ConnectModeManual, AutoInstall: true, - SSHOpts: &SSHOpts{Local: true}, + SSHOpts: &SSHOpts{Local: true, IsSudo: true}, RemoteOpts: &RemoteOptsType{Color: "red"}, Local: true, }