1Panel/backend/utils/websocket/process_data.go
2023-07-10 14:43:12 +00:00

398 lines
9.3 KiB
Go

package websocket
import (
"encoding/json"
"fmt"
"github.com/1Panel-dev/1Panel/backend/global"
"github.com/1Panel-dev/1Panel/backend/utils/files"
"github.com/shirou/gopsutil/v3/host"
"github.com/shirou/gopsutil/v3/net"
"github.com/shirou/gopsutil/v3/process"
"sort"
"strings"
"sync"
"time"
)
type WsInput struct {
Type string `json:"type"`
DownloadProgress
PsProcessConfig
SSHSessionConfig
NetConfig
}
type DownloadProgress struct {
Keys []string `json:"keys"`
}
type PsProcessConfig struct {
Pid int32 `json:"pid"`
Name string `json:"name"`
Username string `json:"username"`
}
type SSHSessionConfig struct {
LoginUser string `json:"loginUser"`
LoginIP string `json:"loginIP"`
}
type NetConfig struct {
Port uint32 `json:"port"`
ProcessName string `json:"processName"`
ProcessID int32 `json:"processID"`
}
type PsProcessData struct {
PID int32 `json:"PID"`
Name string `json:"name"`
PPID int32 `json:"PPID"`
Username string `json:"username"`
Status string `json:"status"`
StartTime string `json:"startTime"`
NumThreads int32 `json:"numThreads"`
NumConnections int `json:"numConnections"`
CpuPercent string `json:"cpuPercent"`
DiskRead string `json:"diskRead"`
DiskWrite string `json:"diskWrite"`
CmdLine string `json:"cmdLine"`
Rss string `json:"rss"`
VMS string `json:"vms"`
HWM string `json:"hwm"`
Data string `json:"data"`
Stack string `json:"stack"`
Locked string `json:"locked"`
Swap string `json:"swap"`
CpuValue float64 `json:"cpuValue"`
RssValue uint64 `json:"rssValue"`
Envs []string `json:"envs"`
OpenFiles []process.OpenFilesStat `json:"openFiles"`
Connects []processConnect `json:"connects"`
}
type processConnect struct {
Type string `json:"type"`
Status string `json:"status"`
Laddr net.Addr `json:"localaddr"`
Raddr net.Addr `json:"remoteaddr"`
PID int32 `json:"PID"`
Name string `json:"name"`
}
type ProcessConnects []processConnect
func (p ProcessConnects) Len() int {
return len(p)
}
func (p ProcessConnects) Less(i, j int) bool {
return p[i].PID < p[j].PID
}
func (p ProcessConnects) Swap(i, j int) {
p[i], p[j] = p[j], p[i]
}
type sshSession struct {
Username string `json:"username"`
PID int32 `json:"PID"`
Terminal string `json:"terminal"`
Host string `json:"host"`
LoginTime string `json:"loginTime"`
}
func ProcessData(c *Client, inputMsg []byte) {
wsInput := &WsInput{}
err := json.Unmarshal(inputMsg, wsInput)
if err != nil {
global.LOG.Errorf("unmarshal wsInput error,err %s", err.Error())
return
}
switch wsInput.Type {
case "wget":
res, err := getDownloadProcess(wsInput.DownloadProgress)
if err != nil {
return
}
c.Msg <- res
case "ps":
res, err := getProcessData(wsInput.PsProcessConfig)
if err != nil {
return
}
c.Msg <- res
case "ssh":
res, err := getSSHSessions(wsInput.SSHSessionConfig)
if err != nil {
return
}
c.Msg <- res
case "net":
res, err := getNetConnections(wsInput.NetConfig)
if err != nil {
return
}
c.Msg <- res
}
}
func getDownloadProcess(progress DownloadProgress) (res []byte, err error) {
var result []files.Process
for _, k := range progress.Keys {
value, err := global.CACHE.Get(k)
if err != nil {
global.LOG.Errorf("get cache error,err %s", err.Error())
return nil, err
}
downloadProcess := &files.Process{}
_ = json.Unmarshal(value, downloadProcess)
result = append(result, *downloadProcess)
}
res, err = json.Marshal(result)
return
}
const (
b = uint64(1)
kb = 1024 * b
mb = 1024 * kb
gb = 1024 * mb
)
func formatBytes(bytes uint64) string {
switch {
case bytes < kb:
return fmt.Sprintf("%dB", bytes)
case bytes < mb:
return fmt.Sprintf("%.2fKB", float64(bytes)/float64(kb))
case bytes < gb:
return fmt.Sprintf("%.2fMB", float64(bytes)/float64(mb))
default:
return fmt.Sprintf("%.2fGB", float64(bytes)/float64(gb))
}
}
func getProcessData(processConfig PsProcessConfig) (res []byte, err error) {
var processes []*process.Process
processes, err = process.Processes()
if err != nil {
return
}
var (
result []PsProcessData
resultMutex sync.Mutex
wg sync.WaitGroup
numWorkers = 4
)
handleData := func(proc *process.Process) {
procData := PsProcessData{
PID: proc.Pid,
}
if processConfig.Pid > 0 && processConfig.Pid != proc.Pid {
return
}
if procName, err := proc.Name(); err == nil {
procData.Name = procName
} else {
procData.Name = "<UNKNOWN>"
}
if processConfig.Name != "" && !strings.Contains(procData.Name, processConfig.Name) {
return
}
if username, err := proc.Username(); err == nil {
procData.Username = username
}
if processConfig.Username != "" && !strings.Contains(procData.Username, processConfig.Username) {
return
}
procData.PPID, _ = proc.Ppid()
statusArray, _ := proc.Status()
if len(statusArray) > 0 {
procData.Status = strings.Join(statusArray, ",")
}
createTime, procErr := proc.CreateTime()
if procErr == nil {
t := time.Unix(createTime/1000, 0)
procData.StartTime = t.Format("2006-1-2 15:04:05")
}
procData.NumThreads, _ = proc.NumThreads()
connections, procErr := proc.Connections()
if procErr == nil {
procData.NumConnections = len(connections)
for _, conn := range connections {
if conn.Laddr.IP != "" || conn.Raddr.IP != "" {
procData.Connects = append(procData.Connects, processConnect{
Status: conn.Status,
Laddr: conn.Laddr,
Raddr: conn.Raddr,
})
}
}
}
procData.CpuValue, _ = proc.CPUPercent()
procData.CpuPercent = fmt.Sprintf("%.2f", procData.CpuValue) + "%"
menInfo, procErr := proc.MemoryInfo()
if procErr == nil {
procData.Rss = formatBytes(menInfo.RSS)
procData.RssValue = menInfo.RSS
procData.Data = formatBytes(menInfo.Data)
procData.VMS = formatBytes(menInfo.VMS)
procData.HWM = formatBytes(menInfo.HWM)
procData.Stack = formatBytes(menInfo.Stack)
procData.Locked = formatBytes(menInfo.Locked)
procData.Swap = formatBytes(menInfo.Swap)
} else {
procData.Rss = "--"
procData.Data = "--"
procData.VMS = "--"
procData.HWM = "--"
procData.Stack = "--"
procData.Locked = "--"
procData.Swap = "--"
procData.RssValue = 0
}
ioStat, procErr := proc.IOCounters()
if procErr == nil {
procData.DiskWrite = formatBytes(ioStat.WriteBytes)
procData.DiskRead = formatBytes(ioStat.ReadBytes)
} else {
procData.DiskWrite = "--"
procData.DiskRead = "--"
}
procData.CmdLine, _ = proc.Cmdline()
procData.OpenFiles, _ = proc.OpenFiles()
procData.Envs, _ = proc.Environ()
resultMutex.Lock()
result = append(result, procData)
resultMutex.Unlock()
}
chunkSize := (len(processes) + numWorkers - 1) / numWorkers
for i := 0; i < numWorkers; i++ {
wg.Add(1)
start := i * chunkSize
end := (i + 1) * chunkSize
if end > len(processes) {
end = len(processes)
}
go func(start, end int) {
defer wg.Done()
for j := start; j < end; j++ {
handleData(processes[j])
}
}(start, end)
}
wg.Wait()
sort.Slice(result, func(i, j int) bool {
return result[i].PID < result[j].PID
})
res, err = json.Marshal(result)
return
}
func getSSHSessions(config SSHSessionConfig) (res []byte, err error) {
var (
result []sshSession
users []host.UserStat
processes []*process.Process
)
processes, err = process.Processes()
if err != nil {
return
}
users, err = host.Users()
if err != nil {
return
}
for _, proc := range processes {
name, _ := proc.Name()
if name != "sshd" || proc.Pid == 0 {
continue
}
connections, _ := proc.Connections()
for _, conn := range connections {
for _, user := range users {
if user.Host == "" {
continue
}
if conn.Raddr.IP == user.Host {
if config.LoginUser != "" && !strings.Contains(user.User, config.LoginUser) {
continue
}
if config.LoginIP != "" && !strings.Contains(user.Host, config.LoginIP) {
continue
}
if terminal, err := proc.Cmdline(); err == nil {
if strings.Contains(terminal, user.Terminal) {
session := sshSession{
Username: user.User,
Host: user.Host,
Terminal: user.Terminal,
PID: proc.Pid,
}
t := time.Unix(int64(user.Started), 0)
session.LoginTime = t.Format("2006-1-2 15:04:05")
result = append(result, session)
}
}
}
}
}
}
res, err = json.Marshal(result)
return
}
var netTypes = [...]string{"tcp", "udp"}
func getNetConnections(config NetConfig) (res []byte, err error) {
var (
result []processConnect
proc *process.Process
)
for _, netType := range netTypes {
connections, _ := net.Connections(netType)
if err == nil {
for _, conn := range connections {
if config.ProcessID > 0 && config.ProcessID != conn.Pid {
continue
}
proc, err = process.NewProcess(conn.Pid)
if err == nil {
name, _ := proc.Name()
if name != "" && config.ProcessName != "" && !strings.Contains(name, config.ProcessName) {
continue
}
if config.Port > 0 && config.Port != conn.Laddr.Port && config.Port != conn.Raddr.Port {
continue
}
result = append(result, processConnect{
Type: netType,
Status: conn.Status,
Laddr: conn.Laddr,
Raddr: conn.Raddr,
PID: conn.Pid,
Name: name,
})
}
}
}
}
res, err = json.Marshal(result)
return
}