Refactors bitness and improves ps list

chennqqi · Oct 2, 2020 · 3aba989 · 3aba989
1 parent 315fdb6
commit 3aba989
Show file tree

Hide file tree

Showing 13 changed files with 152 additions and 55 deletions.
diff --git a/app/app.go b/app/app.go
@@ -225,6 +225,13 @@ func RunApp(args []string) {
 				Aliases: []string{"ps", "lsproc"},
 				Usage:   "lists all running processes",
 				Action:  listProcesses,
+				Flags: []cli.Flag{
+					&cli.BoolFlag{
+						Name:    "verbose",
+						Aliases: []string{"v"},
+						Usage:   "output errors if any are encountered",
+					},
+				},
 			},
 			&cli.Command{
 				Name:      "list-process-memory",

diff --git a/app/listProcs.go b/app/listProcs.go
@@ -26,7 +26,9 @@ func listProcesses(c *cli.Context) error {
 	maxPidlen := 0
 	maxNamelen := 0
 	maxUserlen := 0
+	errorsOutput := false
 	for i, pid := range pids {
+		// Default info in case of errors
 		info := &procIO.ProcessInfo{
 			PID:              pid,
 			ExecutablePath:   "ERROR",
@@ -37,11 +39,20 @@ func listProcesses(c *cli.Context) error {
 		}
 
 		proc, err := procIO.OpenProcess(pid)
-		if err == nil {
+		if err != nil {
+			err = errors.Newf("could not open process %d, reason: %w", pid, err)
+		} else {
 			tmp, err := proc.Info()
-			if err == nil {
+			if err != nil {
+				err = errors.Newf("could not query info of process %d, reason: %w", pid, err)
+			} else {
 				info = tmp
 			}
+			proc.Close()
+		}
+		if c.Bool("verbose") && err != nil {
+			errorsOutput = true
+			fmt.Println(err)
 		}
 
 		procInfos[i] = info
@@ -63,12 +74,17 @@ func listProcesses(c *cli.Context) error {
 		maxPidlen = 5
 	}
 
-	headerFmt := fmt.Sprintf("%%%ds %%-%ds %%-%ds\n", maxPidlen, maxNamelen, maxUserlen)
-	rowFmt := fmt.Sprintf("%%%dd %%-%ds %%-%ds\n", maxPidlen, maxNamelen, maxUserlen)
-	fmt.Printf(headerFmt, "PID", "Name", "User")
-	fmt.Println(strings.Repeat("-", maxPidlen) + "+" + strings.Repeat("-", maxNamelen) + "+" + strings.Repeat("-", maxUserlen))
+	if errorsOutput {
+		// Extra empty line for readability
+		fmt.Println()
+	}
+
+	headerFmt := fmt.Sprintf("%%%ds %%3v %%-%ds %%-%ds\n", maxPidlen, maxNamelen, maxUserlen)
+	rowFmt := fmt.Sprintf("%%%dd %%3v %%-%ds %%-%ds\n", maxPidlen, maxNamelen, maxUserlen)
+	fmt.Printf(headerFmt, "PID", "Bit", "Name", "User")
+	fmt.Println(strings.Repeat("-", maxPidlen) + "+" + strings.Repeat("-", 3) + "+" + strings.Repeat("-", maxNamelen) + "+" + strings.Repeat("-", maxUserlen))
 	for _, info := range procInfos {
-		fmt.Printf(rowFmt, info.PID, filepath.Base(info.ExecutablePath), info.Username)
+		fmt.Printf(rowFmt, info.PID, info.Bitness.Short(), filepath.Base(info.ExecutablePath), info.Username)
 	}
 
 	return nil

diff --git a/arch/arch.go b/arch/arch.go
@@ -0,0 +1,19 @@
+package arch
+
+type T int
+
+const (
+	Invalid T = iota
+	AMD64
+	I386
+)
+
+var bitness = map[T]Bitness{
+	Invalid: BitnessInvalid,
+	AMD64:   Bitness64Bit,
+	I386:    Bitness32Bit,
+}
+
+func (t T) Bitness() Bitness {
+	return bitness[t]
+}
diff --git a/arch/bitness.go b/arch/bitness.go
@@ -0,0 +1,21 @@
+//go:generate go-enum -f=$GOFILE --marshal
+package arch
+
+/*
+ENUM(
+Invalid
+32Bit = 32
+64Bit = 64
+)
+*/
+type Bitness int
+
+var bitnessShortNames = map[Bitness]string{
+	BitnessInvalid: "??",
+	Bitness64Bit:   "64",
+	Bitness32Bit:   "32",
+}
+
+func (b Bitness) Short() string {
+	return bitnessShortNames[b]
+}
diff --git a/system/info_enum.go → arch/bitness_enum.go b/system/info_enum.go → arch/bitness_enum.go
diff --git a/arch/get.go b/arch/get.go
@@ -0,0 +1,5 @@
+package arch
+
+func Native() T {
+	return get()
+}
diff --git a/arch/get_386.go b/arch/get_386.go
@@ -0,0 +1,5 @@
+package arch
+
+func get() T {
+	return I386
+}
diff --git a/arch/get_amd64.go b/arch/get_amd64.go
@@ -0,0 +1,5 @@
+package arch
+
+func get() T {
+	return AMD64
+}
diff --git a/procIO/process.go b/procIO/process.go
@@ -6,6 +6,7 @@ import (
 	"encoding/hex"
 	"errors"
 	"fmt"
+	"fraunhofer/fkie/yapscan/arch"
 	"io"
 	"os"
 )
@@ -15,6 +16,7 @@ var ErrProcIsParent = errors.New("not supported on parent")
 
 type ProcessInfo struct {
 	PID              int                  `json:"pid"`
+	Bitness          arch.Bitness         `json:"bitness"`
 	ExecutablePath   string               `json:"executablePath"`
 	ExecutableMD5    string               `json:"executableMD5"`
 	ExecutableSHA256 string               `json:"executableSHA256"`

diff --git a/procIO/process_windows.go b/procIO/process_windows.go
@@ -1,16 +1,40 @@
 package procIO
 
 import (
+	"fraunhofer/fkie/yapscan/arch"
 	"fraunhofer/fkie/yapscan/procIO/customWin32"
 	"os"
 	"syscall"
 
+	"golang.org/x/sys/windows"
+
 	"github.com/targodan/go-errors"
 
 	"github.com/0xrawsec/golang-win32/win32"
 	"github.com/0xrawsec/golang-win32/win32/kernel32"
 )
 
+var specialPIDs = map[int]*ProcessInfo{
+	0: &ProcessInfo{
+		PID:              0,
+		Bitness:          arch.Native().Bitness(),
+		ExecutablePath:   "IdleProcess",
+		ExecutableMD5:    "",
+		ExecutableSHA256: "",
+		Username:         "System",
+		MemorySegments:   []*MemorySegmentInfo{},
+	},
+	4: &ProcessInfo{
+		PID:              4,
+		Bitness:          arch.Native().Bitness(),
+		ExecutablePath:   "PsInitialSystemProcess",
+		ExecutableMD5:    "",
+		ExecutableSHA256: "",
+		Username:         "System",
+		MemorySegments:   []*MemorySegmentInfo{},
+	},
+}
+
 func GetRunningPIDs() ([]int, error) {
 	snap, err := kernel32.CreateToolhelp32Snapshot(kernel32.TH32CS_SNAPPROCESS, 0)
 	if err != nil {
@@ -47,6 +71,11 @@ type processWindows struct {
 }
 
 func open(pid int) (Process, error) {
+	if pid <= 4 {
+		// We'll create special processes without handle, so the info can at least be retreived
+		return &processWindows{pid: pid, procHandle: 0}, nil
+	}
+
 	handle, err := kernel32.OpenProcess(
 		kernel32.PROCESS_VM_READ|kernel32.PROCESS_QUERY_INFORMATION|kernel32.PROCESS_SUSPEND_RESUME,
 		win32.FALSE,
@@ -64,6 +93,11 @@ func (p *processWindows) PID() int {
 }
 
 func (p *processWindows) Info() (*ProcessInfo, error) {
+	special, ok := specialPIDs[p.pid]
+	if ok {
+		return special, nil
+	}
+
 	var tmpErr, err error
 	info := &ProcessInfo{
 		PID: p.pid,
@@ -74,6 +108,19 @@ func (p *processWindows) Info() (*ProcessInfo, error) {
 		err = errors.NewMultiError(err, errors.Errorf("could not retrieve memory segments info, reason: %w", tmpErr))
 	}
 
+	var isWow64 bool
+	err = windows.IsWow64Process(windows.Handle(p.procHandle), &isWow64)
+	if tmpErr != nil {
+		err = errors.NewMultiError(err, errors.Errorf("could not determine process bitness, reason: %w", tmpErr))
+	}
+	// Note: This is good for windows on x86 and x86_64.
+	// Docs: https://docs.microsoft.com/en-us/windows/win32/api/wow64apiset/nf-wow64apiset-iswow64process?redirectedfrom=MSDN
+	if isWow64 {
+		info.Bitness = arch.Bitness32Bit
+	} else {
+		info.Bitness = arch.Bitness64Bit
+	}
+
 	info.ExecutablePath, tmpErr = kernel32.GetModuleFilenameExW(p.procHandle, 0)
 	if tmpErr != nil {
 		err = errors.NewMultiError(err, errors.Errorf("could not retrieve executable path, reason: %w", tmpErr))

diff --git a/system/info.go b/system/info.go
@@ -1,26 +1,18 @@
-//go:generate go-enum -f=$GOFILE --marshal
 package system
 
 import (
+	"fraunhofer/fkie/yapscan/arch"
 	"net"
 	"os"
 
 	"github.com/targodan/go-errors"
 )
 
-/*
-ENUM(
-32Bit
-64Bit
-)
-*/
-type Bitness int
-
 type Info struct {
 	OSName    string   `json:"osName"`
 	OSVersion string   `json:"osVersion"`
 	OSFlavour string   `json:"osFlavour"`
-	OSBitness Bitness  `json:"osBitness"`
+	OSArch    arch.T   `json:"osArch"`
 	Hostname  string   `json:"hostname"`
 	IPs       []string `json:"ips"`
 }
@@ -32,7 +24,8 @@ func GetInfo() (*Info, error) {
 		var err error
 
 		info = new(Info)
-		info.OSName, info.OSVersion, info.OSFlavour, info.OSBitness, err = getOSInfo()
+		info.OSArch = arch.Native()
+		info.OSName, info.OSVersion, info.OSFlavour, err = getOSInfo()
 		if err != nil {
 			err = errors.Errorf("could not determine OS info, reason: %w", err)
 			return info, err

diff --git a/system/info_linux.go b/system/info_linux.go
@@ -3,11 +3,9 @@ package system
 import (
 	"os/exec"
 	"strings"
-
-	"github.com/targodan/go-errors"
 )
 
-func getOSInfo() (name, version, flavour string, bitness Bitness, err error) {
+func getOSInfo() (name, version, flavour string, err error) {
 	var buf []byte
 
 	cmd := exec.Command("uname", "-s")
@@ -31,22 +29,6 @@ func getOSInfo() (name, version, flavour string, bitness Bitness, err error) {
 	}
 	arch := strings.TrimSpace(string(buf))
 
-	switch arch {
-	case "amd64":
-		fallthrough
-	case "x86_64":
-		bitness = Bitness64Bit
-
-	case "i686":
-		fallthrough
-	case "x86":
-		bitness = Bitness32Bit
-
-	default:
-		err = errors.Errorf("unknown architecture \"%s\"", arch)
-		return
-	}
-
 	cmd = exec.Command("uname", "-o")
 	buf, err = cmd.Output()
 	if err != nil {

diff --git a/system/info_windows.go b/system/info_windows.go
@@ -2,22 +2,13 @@ package system
 
 import (
 	"encoding/csv"
-	"os"
 	"os/exec"
 	"strings"
 
 	"github.com/targodan/go-errors"
 )
 
-func getOSInfo() (name, version, flavour string, bitness Bitness, err error) {
-	_, exists := os.LookupEnv("ProgramFiles(x86)")
-	if exists {
-		// Is 64bit OS
-		bitness = Bitness64Bit
-	} else {
-		// Is not 64bit OS, assume 32bit
-		bitness = Bitness32Bit
-	}
+func getOSInfo() (name, version, flavour string, err error) {
 	cmd := exec.Command("systeminfo", "/FO", "CSV")
 	buf, err := cmd.Output()
 	if err != nil {