staros/process_unix.go

// +build linux darwin

package staros

import (
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"
	"time"
)

//FindProcessByName 通过进程名来查询应用信息
func FindProcessByName(name string) (datas []Process, err error) {
	return FindProcess(func(in Process) bool {
		if name == in.Name {
			return true
		}
		return false
	})
}

// FindProcess 通过进程信息来查询应用信息
func FindProcess(compare func(Process) bool) (datas []Process, err error) {
	var name, main string
	var mainb []byte
	var netErr error
	var netInfo []NetConn
	paths, err := ioutil.ReadDir("/proc")
	if err != nil {
		return
	}
	netInfo, netErr = NetConnections(false, "")
	appendNetInfo := func(p *Process) {
		if netErr != nil {
			p.netErr = netErr
			return
		}
		fds, err := ioutil.ReadDir("/proc/" + strconv.Itoa(int(p.Pid)) + "/fd")
		if err != nil && Exists("/proc/"+strconv.Itoa(int(p.Pid))+"/fd") {
			p.netErr = err
			return
		}
		for _, fd := range fds {
			socket, err := os.Readlink("/proc/" + strconv.Itoa(int(p.Pid)) + "/fd/" + fd.Name())
			if err != nil {
				p.netErr = err
				return
			}
			start := strings.Index(socket, "[")
			if start < 0 {
				continue
			}
			sid := socket[start+1 : len(socket)-1]
			for _, v := range netInfo {
				if v.Inode == sid {
					v.Pid = p.Pid
					v.Process = p
					p.netConn = append(p.netConn, v)
				}
			}
		}
	}
	for _, v := range paths {
		if v.IsDir() && Exists("/proc/"+v.Name()+"/comm") {
			name, err = readAsString("/proc/" + v.Name() + "/comm")
			if err != nil {
				continue
			}
			var tmp Process
			tmp.LocalPath, err = os.Readlink("/proc/" + v.Name() + "/exe")
			tmp.Path = tmp.LocalPath
			tmp.LocalPath = filepath.Dir(tmp.LocalPath)
			tmp.ExecPath, err = os.Readlink("/proc/" + v.Name() + "/cwd")
			tmp.Name = strings.TrimSpace(name)
			main, err = readAsString("/proc/" + v.Name() + "/status")
			if err != nil {
				tmp.Err = err
				if compare(tmp) {
					appendNetInfo(&tmp)
					datas = append(datas, tmp)
					continue
				}
			} else {
				data := splitBy(main, ":")
				tmp.Pid, _ = strconv.ParseInt(data["Pid"], 10, 64)
				tmp.PPid, _ = strconv.ParseInt(data["PPid"], 10, 64)
				tmp.TPid, _ = strconv.ParseInt(data["TracerPid"], 10, 64)
				uids := splitBySpace(data["Uid"])
				gids := splitBySpace(data["Gid"])
				tmp.RUID, _ = strconv.Atoi(uids[0])
				tmp.EUID, _ = strconv.Atoi(uids[1])
				tmp.RGID, _ = strconv.Atoi(gids[0])
				tmp.EGID, _ = strconv.Atoi(gids[1])
				tmp.VmPeak, _ = strconv.ParseInt(splitBySpace(data["VmPeak"])[0], 10, 64)
				tmp.VmSize, _ = strconv.ParseInt(splitBySpace(data["VmSize"])[0], 10, 64)
				tmp.VmHWM, _ = strconv.ParseInt(splitBySpace(data["VmHWM"])[0], 10, 64)
				tmp.VmRSS, _ = strconv.ParseInt(splitBySpace(data["VmRSS"])[0], 10, 64)
				tmp.VmLck, _ = strconv.ParseInt(splitBySpace(data["VmLck"])[0], 10, 64)
				tmp.VmData, _ = strconv.ParseInt(splitBySpace(data["VmData"])[0], 10, 64)
				tmp.VmLck *= 1024
				tmp.VmData *= 1024
				tmp.VmPeak *= 1024
				tmp.VmSize *= 1024
				tmp.VmHWM *= 1024
				tmp.VmRSS *= 1024
			}
			mainb, err = ioutil.ReadFile("/proc/" + v.Name() + "/cmdline")
			if err != nil {
				tmp.Err = err
				if compare(tmp) {
					appendNetInfo(&tmp)
					datas = append(datas, tmp)
					continue
				}
			} else {
				args := bytes.Split(mainb, []byte{0})
				for _, v := range args {
					tmp.Args = append(tmp.Args, string(v))
				}
			}
			mainb, err = ioutil.ReadFile("/proc/" + v.Name() + "/environ")
			if err != nil {
				tmp.Err = err
				if compare(tmp) {
					appendNetInfo(&tmp)
					datas = append(datas, tmp)
					continue
				}
			} else {
				args := bytes.Split(mainb, []byte{0})
				for _, v := range args {
					tmp.Env = append(tmp.Env, string(v))
				}
			}

			main, err = readAsString("/proc/" + v.Name() + "/stat")
			if err != nil {
				tmp.Err = err
				if compare(tmp) {
					appendNetInfo(&tmp)
					datas = append(datas, tmp)
					continue
				}
			} else {
				times := splitBySpace(main)
				uptime, _ := strconv.ParseInt(strings.TrimSpace(times[21]), 10, 64)
				tmp.Uptime = time.Unix(StartTime().Unix()+uptime/100, int64((float64(uptime)/100-float64(uptime/100))*1000000000))
			}
			if compare(tmp) {
				appendNetInfo(&tmp)
				datas = append(datas, tmp)
			}
		}
	}
	return
}

// FindProcessByPid 通过Pid来查询应用信息
func FindProcessByPid(pid int64) (datas Process, err error) {
	var name, main string
	var mainb []byte
	if !Exists("/proc/" + fmt.Sprint(pid) + "/comm") {
		err = errors.New("Not Found")
		return
	}
	netInfo, netErr := NetConnections(false, "")
	appendNetInfo := func(p *Process) {
		if netErr != nil {
			p.netErr = netErr
			return
		}
		fds, err := ioutil.ReadDir("/proc/" + strconv.Itoa(int(p.Pid)) + "/fd")
		if err != nil && Exists("/proc/"+strconv.Itoa(int(p.Pid))+"/fd") {
			p.netErr = err
			return
		}
		for _, fd := range fds {
			socket, err := os.Readlink("/proc/" + strconv.Itoa(int(p.Pid)) + "/fd/" + fd.Name())
			if err != nil {
				p.netErr = err
				return
			}
			start := strings.Index(socket, "[")
			if start < 0 {
				continue
			}
			sid := socket[start+1 : len(socket)-1]
			for _, v := range netInfo {
				if v.Inode == sid {
					v.Pid = p.Pid
					v.Process = p
					p.netConn = append(p.netConn, v)
				}
			}
		}
	}

	name, err = readAsString("/proc/" + fmt.Sprint(pid) + "/comm")
	if err != nil {
		return
	}
	main, err = readAsString("/proc/" + fmt.Sprint(pid) + "/status")
	if err != nil {
		return
	}
	data := splitBy(main, ":")
	datas.Name = strings.TrimSpace(name)
	datas.Pid, _ = strconv.ParseInt(data["Pid"], 10, 64)
	datas.PPid, _ = strconv.ParseInt(data["PPid"], 10, 64)
	datas.TPid, _ = strconv.ParseInt(data["TracerPid"], 10, 64)
	uids := splitBySpace(data["Uid"])
	gids := splitBySpace(data["Gid"])
	datas.RUID, _ = strconv.Atoi(uids[0])
	datas.EUID, _ = strconv.Atoi(uids[1])
	datas.RGID, _ = strconv.Atoi(gids[0])
	datas.EGID, _ = strconv.Atoi(gids[1])
	datas.VmPeak, _ = strconv.ParseInt(splitBySpace(data["VmPeak"])[0], 10, 64)
	datas.VmSize, _ = strconv.ParseInt(splitBySpace(data["VmSize"])[0], 10, 64)
	datas.VmHWM, _ = strconv.ParseInt(splitBySpace(data["VmHWM"])[0], 10, 64)
	datas.VmRSS, _ = strconv.ParseInt(splitBySpace(data["VmRSS"])[0], 10, 64)
	datas.VmLck, _ = strconv.ParseInt(splitBySpace(data["VmLck"])[0], 10, 64)
	datas.VmData, _ = strconv.ParseInt(splitBySpace(data["VmData"])[0], 10, 64)
	datas.VmLck *= 1024
	datas.VmData *= 1024
	datas.VmPeak *= 1024
	datas.VmSize *= 1024
	datas.VmHWM *= 1024
	datas.VmRSS *= 1024
	appendNetInfo(&datas)
	mainb, err = ioutil.ReadFile("/proc/" + fmt.Sprint(pid) + "/cmdline")
	if err != nil {
		datas.Err = err
		err = nil
	} else {
		args := bytes.Split(mainb, []byte{0})
		for _, v := range args {
			datas.Args = append(datas.Args, string(v))
		}
	}

	mainb, err = ioutil.ReadFile("/proc/" + fmt.Sprint(pid) + "/environ")
	if err != nil {
		datas.Err = err
		err = nil
	} else {
		args := bytes.Split(mainb, []byte{0})
		for _, v := range args {
			datas.Env = append(datas.Env, string(v))
		}
	}

	datas.LocalPath, err = os.Readlink("/proc/" + fmt.Sprint(pid) + "/exe")
	datas.Path = datas.LocalPath
	datas.LocalPath = filepath.Dir(datas.LocalPath)
	datas.ExecPath, err = os.Readlink("/proc/" + fmt.Sprint(pid) + "/cwd")
	main, err = readAsString("/proc/" + fmt.Sprint(pid) + "/stat")
	if err != nil {
		return
	}
	times := splitBySpace(main)
	uptime, _ := strconv.ParseInt(strings.TrimSpace(times[21]), 10, 64)
	datas.Uptime = time.Unix(StartTime().Unix()+uptime/100, int64((float64(uptime)/100-float64(uptime/100))*1000000000))
	return
}

func Daemon(path string, args ...string) (int, error) {
	cmd := exec.Command(path, args...)
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Setsid: true,
	}
	if err := cmd.Start(); err != nil {
		return -1, err
	}
	pid := cmd.Process.Pid
	err := cmd.Process.Release()
	return pid, err
}

func DaemonWithUser(uid, gid uint32, groups []uint32, path string, args ...string) (int, error) {
	cmd := exec.Command(path, args...)
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Credential: &syscall.Credential{
			Uid:    uid,
			Gid:    gid,
			Groups: groups,
		},
		Setsid: true,
	}
	if err := cmd.Start(); err != nil {
		return -1, err
	}
	pid := cmd.Process.Pid
	err := cmd.Process.Release()
	return pid, err
}

func (starcli *StarCmd) SetRunUser(uid, gid uint32, groups []uint32) {
	starcli.CMD.SysProcAttr = &syscall.SysProcAttr{
		Credential: &syscall.Credential{
			Uid:    uid,
			Gid:    gid,
			Groups: groups,
		},
		Setsid: true,
	}
}

func (starcli *StarCmd) Release() error {
	if starcli.CMD.SysProcAttr == nil {
		starcli.CMD.SysProcAttr = &syscall.SysProcAttr{
			Setsid: true,
		}
	} else {
		if !starcli.CMD.SysProcAttr.Setsid {
			starcli.CMD.SysProcAttr.Setsid = true
		}
	}
	if !starcli.IsRunning() {
		if err := starcli.CMD.Start(); err != nil {
			return err
		}
	}
	time.Sleep(time.Millisecond * 10)
	return starcli.CMD.Process.Release()
}

func (starcli *StarCmd) SetKeepCaps() error {
	_, _, err := syscall.RawSyscall(157 /*SYS PRCTL */, 0x8 /*PR SET KEEPCAPS*/, 1, 0)
	if 0 != err {
		return err
	}
	return nil
}

func SetKeepCaps() error {
	_, _, err := syscall.RawSyscall(157 /*SYS PRCTL */, 0x8 /*PR SET KEEPCAPS*/, 1, 0)
	if 0 != err {
		return err
	}
	return nil
}