Merge pull request shirou#1163 from secDre4mer/master

feat: Add support for Cwd() on Windows
kazdevl · Oct 28, 2021 · 9928258 · 9928258
2 parents fb65e18 + a4679b7
commit 9928258
Show file tree

Hide file tree

Showing 4 changed files with 220 additions and 26 deletions.
diff --git a/process/process_test.go b/process/process_test.go
@@ -720,6 +720,22 @@ func Test_IsRunning(t *testing.T) {
 	}
 }
 
+func Test_Process_Cwd(t *testing.T) {
+	myPid := os.Getpid()
+	currentWorkingDirectory, _ := os.Getwd()
+
+	process, _ := NewProcess(int32(myPid))
+	pidCwd, err := process.Cwd()
+	skipIfNotImplementedErr(t, err)
+	if err != nil {
+		t.Fatalf("getting cwd error %v", err)
+	}
+	pidCwd = strings.TrimSuffix(pidCwd, string(os.PathSeparator))
+	assert.Equal(t, currentWorkingDirectory, pidCwd)
+
+	t.Log(pidCwd)
+}
+
 func Test_Process_Environ(t *testing.T) {
 	tmpdir, err := ioutil.TempDir("", "")
 	if err != nil {
@@ -811,6 +827,23 @@ func Test_AllProcesses_environ(t *testing.T) {
 	}
 }
 
+func Test_AllProcesses_Cwd(t *testing.T) {
+	procs, err := Processes()
+	skipIfNotImplementedErr(t, err)
+	if err != nil {
+		t.Fatalf("getting processes error %v", err)
+	}
+	for _, proc := range procs {
+		exeName, _ := proc.Exe()
+		cwd, err := proc.Cwd()
+		if err != nil {
+			cwd = "Error: " + err.Error()
+		}
+
+		t.Logf("Process #%v: Name: %v / Current Working Directory: %s\n", proc.Pid, exeName, cwd)
+	}
+}
+
 func BenchmarkNewProcess(b *testing.B) {
 	checkPid := os.Getpid()
 	for i := 0; i < b.N; i++ {

diff --git a/process/process_windows.go b/process/process_windows.go
@@ -137,30 +137,54 @@ type processEnvironmentBlock64 struct {
 }
 
 type rtlUserProcessParameters32 struct {
-	Reserved1           [16]uint8
-	Reserved2           [10]uint32
+	Reserved1 [16]uint8
+	ConsoleHandle uint32
+	ConsoleFlags uint32
+	StdInputHandle uint32
+	StdOutputHandle uint32
+	StdErrorHandle uint32
+	CurrentDirectoryPathNameLength uint16
+	_ uint16 // Max Length
+	CurrentDirectoryPathAddress uint32
+	CurrentDirectoryHandle uint32
+	DllPathNameLength uint16
+	_ uint16 // Max Length
+	DllPathAddress uint32
 	ImagePathNameLength uint16
-	_                   uint16
-	ImagePathAddress    uint32
-	CommandLineLength   uint16
-	_                   uint16
-	CommandLineAddress  uint32
-	EnvironmentAddress  uint32
+	_ uint16 // Max Length
+	ImagePathAddress uint32
+	CommandLineLength uint16
+	_ uint16 // Max Length
+	CommandLineAddress uint32
+	EnvironmentAddress uint32
 	// More fields which we don't use so far
 }
 
 type rtlUserProcessParameters64 struct {
-	Reserved1           [16]uint8
-	Reserved2           [10]uint64
+	Reserved1 [16]uint8
+	ConsoleHandle uint64
+	ConsoleFlags uint64
+	StdInputHandle uint64
+	StdOutputHandle uint64
+	StdErrorHandle uint64
+	CurrentDirectoryPathNameLength uint16
+	_ uint16 // Max Length
+	_ uint32 // Padding
+	CurrentDirectoryPathAddress uint64
+	CurrentDirectoryHandle uint64
+	DllPathNameLength uint16
+	_ uint16 // Max Length
+	_ uint32 // Padding
+	DllPathAddress uint64
 	ImagePathNameLength uint16
-	_                   uint16 // Max Length
-	_                   uint32 // Padding
-	ImagePathAddress    uint64
-	CommandLineLength   uint16
-	_                   uint16 // Max Length
-	_                   uint32 // Padding
-	CommandLineAddress  uint64
-	EnvironmentAddress  uint64
+	_ uint16 // Max Length
+	_ uint32 // Padding
+	ImagePathAddress uint64
+	CommandLineLength uint16
+	_ uint16 // Max Length
+	_ uint32 // Padding
+	CommandLineAddress uint64
+	EnvironmentAddress uint64
 	// More fields which we don't use so far
 }
 
@@ -377,8 +401,48 @@ func (p *Process) createTimeWithContext(ctx context.Context) (int64, error) {
 	return ru.CreationTime.Nanoseconds() / 1000000, nil
 }
 
-func (p *Process) CwdWithContext(ctx context.Context) (string, error) {
-	return "", common.ErrNotImplementedError
+func (p *Process) CwdWithContext(_ context.Context) (string, error) {
+	h, err := windows.OpenProcess(processQueryInformation|windows.PROCESS_VM_READ, false, uint32(p.Pid))
+	if err == windows.ERROR_ACCESS_DENIED || err == windows.ERROR_INVALID_PARAMETER {
+		return "", nil
+	}
+	if err != nil {
+		return "", err
+	}
+	defer syscall.CloseHandle(syscall.Handle(h))
+
+	procIs32Bits := is32BitProcess(h)
+
+	if procIs32Bits {
+		userProcParams, err := getUserProcessParams32(h)
+		if err != nil {
+			return "", err
+		}
+		if userProcParams.CurrentDirectoryPathNameLength > 0 {
+			cwd := readProcessMemory(syscall.Handle(h), procIs32Bits, uint64(userProcParams.CurrentDirectoryPathAddress), uint(userProcParams.CurrentDirectoryPathNameLength))
+			if len(cwd) != int(userProcParams.CurrentDirectoryPathAddress) {
+				return "", errors.New("cannot read current working directory")
+			}
+
+			return convertUTF16ToString(cwd), nil
+		}
+	} else {
+		userProcParams, err := getUserProcessParams64(h)
+		if err != nil {
+			return "", err
+		}
+		if userProcParams.CurrentDirectoryPathNameLength > 0 {
+			cwd := readProcessMemory(syscall.Handle(h), procIs32Bits, userProcParams.CurrentDirectoryPathAddress, uint(userProcParams.CurrentDirectoryPathNameLength))
+			if len(cwd) != int(userProcParams.CurrentDirectoryPathNameLength) {
+				return "", errors.New("cannot read current working directory")
+			}
+
+			return convertUTF16ToString(cwd), nil
+		}
+	}
+
+	//if we reach here, we have no cwd
+	return "", nil
 }
 
 func (p *Process) ParentWithContext(ctx context.Context) (*Process, error) {

diff --git a/v3/process/process_test.go b/v3/process/process_test.go
@@ -776,6 +776,22 @@ func Test_Process_Environ(t *testing.T) {
 	}
 }
 
+func Test_Process_Cwd(t *testing.T) {
+	myPid := os.Getpid()
+	currentWorkingDirectory, _ := os.Getwd()
+
+	process, _ := NewProcess(int32(myPid))
+	pidCwd, err := process.Cwd()
+	skipIfNotImplementedErr(t, err)
+	if err != nil {
+		t.Fatalf("getting cwd error %v", err)
+	}
+	pidCwd = strings.TrimSuffix(pidCwd, string(os.PathSeparator))
+	assert.Equal(t, currentWorkingDirectory, pidCwd)
+
+	t.Log(pidCwd)
+}
+
 func Test_AllProcesses_cmdLine(t *testing.T) {
 	procs, err := Processes()
 	skipIfNotImplementedErr(t, err)
@@ -813,6 +829,23 @@ func Test_AllProcesses_environ(t *testing.T) {
 	}
 }
 
+func Test_AllProcesses_Cwd(t *testing.T) {
+	procs, err := Processes()
+	skipIfNotImplementedErr(t, err)
+	if err != nil {
+		t.Fatalf("getting processes error %v", err)
+	}
+	for _, proc := range procs {
+		exeName, _ := proc.Exe()
+		cwd, err := proc.Cwd()
+		if err != nil {
+			cwd = "Error: " + err.Error()
+		}
+
+		t.Logf("Process #%v: Name: %v / Current Working Directory: %s\n", proc.Pid, exeName, cwd)
+	}
+}
+
 func BenchmarkNewProcess(b *testing.B) {
 	checkPid := os.Getpid()
 	for i := 0; i < b.N; i++ {

diff --git a/v3/process/process_windows.go b/v3/process/process_windows.go
@@ -125,20 +125,44 @@ type processEnvironmentBlock64 struct {
 
 type rtlUserProcessParameters32 struct {
 	Reserved1 [16]uint8
-	Reserved2 [10]uint32
+	ConsoleHandle uint32
+	ConsoleFlags uint32
+	StdInputHandle uint32
+	StdOutputHandle uint32
+	StdErrorHandle uint32
+	CurrentDirectoryPathNameLength uint16
+	_ uint16 // Max Length
+	CurrentDirectoryPathAddress uint32
+	CurrentDirectoryHandle uint32
+	DllPathNameLength uint16
+	_ uint16 // Max Length
+	DllPathAddress uint32
 	ImagePathNameLength uint16
-	_ uint16
+	_ uint16 // Max Length
 	ImagePathAddress uint32
 	CommandLineLength uint16
-	_ uint16
+	_ uint16 // Max Length
 	CommandLineAddress uint32
 	EnvironmentAddress uint32
 	// More fields which we don't use so far
 }
 
 type rtlUserProcessParameters64 struct {
 	Reserved1 [16]uint8
-	Reserved2 [10]uint64
+	ConsoleHandle uint64
+	ConsoleFlags uint64
+	StdInputHandle uint64
+	StdOutputHandle uint64
+	StdErrorHandle uint64
+	CurrentDirectoryPathNameLength uint16
+	_ uint16 // Max Length
+	_ uint32 // Padding
+	CurrentDirectoryPathAddress uint64
+	CurrentDirectoryHandle uint64
+	DllPathNameLength uint16
+	_ uint16 // Max Length
+	_ uint32 // Padding
+	DllPathAddress uint64
 	ImagePathNameLength uint16
 	_ uint16 // Max Length
 	_ uint32 // Padding
@@ -364,8 +388,48 @@ func (p *Process) createTimeWithContext(ctx context.Context) (int64, error) {
 	return ru.CreationTime.Nanoseconds() / 1000000, nil
 }
 
-func (p *Process) CwdWithContext(ctx context.Context) (string, error) {
-	return "", common.ErrNotImplementedError
+func (p *Process) CwdWithContext(_ context.Context) (string, error) {
+	h, err := windows.OpenProcess(processQueryInformation|windows.PROCESS_VM_READ, false, uint32(p.Pid))
+	if err == windows.ERROR_ACCESS_DENIED || err == windows.ERROR_INVALID_PARAMETER {
+		return "", nil
+	}
+	if err != nil {
+		return "", err
+	}
+	defer syscall.CloseHandle(syscall.Handle(h))
+
+	procIs32Bits := is32BitProcess(h)
+
+	if procIs32Bits {
+		userProcParams, err := getUserProcessParams32(h)
+		if err != nil {
+			return "", err
+		}
+		if userProcParams.CurrentDirectoryPathNameLength > 0 {
+			cwd := readProcessMemory(syscall.Handle(h), procIs32Bits, uint64(userProcParams.CurrentDirectoryPathAddress), uint(userProcParams.CurrentDirectoryPathNameLength))
+			if len(cwd) != int(userProcParams.CurrentDirectoryPathAddress) {
+				return "", errors.New("cannot read current working directory")
+			}
+
+			return convertUTF16ToString(cwd), nil
+		}
+	} else {
+		userProcParams, err := getUserProcessParams64(h)
+		if err != nil {
+			return "", err
+		}
+		if userProcParams.CurrentDirectoryPathNameLength > 0 {
+			cwd := readProcessMemory(syscall.Handle(h), procIs32Bits, userProcParams.CurrentDirectoryPathAddress, uint(userProcParams.CurrentDirectoryPathNameLength))
+			if len(cwd) != int(userProcParams.CurrentDirectoryPathNameLength) {
+				return "", errors.New("cannot read current working directory")
+			}
+
+			return convertUTF16ToString(cwd), nil
+		}
+	}
+
+	//if we reach here, we have no cwd
+	return "", nil
 }
 
 func (p *Process) ParentWithContext(ctx context.Context) (*Process, error) {