From 0d366149b98ccb67e3eb4142b1d6a0c5208d8934 Mon Sep 17 00:00:00 2001 From: Jan De Dobbeleer Date: Sun, 6 Mar 2022 08:41:02 +0100 Subject: [PATCH] feat: query all window titles for app info resolves #1354 --- .vscode/launch.json | 2 +- src/environment/shell.go | 2 +- src/environment/unix.go | 2 +- src/environment/windows.go | 6 +- src/environment/windows_win32.go | 180 +++++++++------------------ src/mock/environment.go | 4 +- src/segments/spotify_windows.go | 20 ++- src/segments/spotify_windows_test.go | 105 ++++++++++------ src/segments/spotify_wsl_test.go | 3 +- 9 files changed, 147 insertions(+), 177 deletions(-) diff --git a/.vscode/launch.json b/.vscode/launch.json index f1355e6b..c089877c 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -8,7 +8,7 @@ "mode": "debug", "program": "${workspaceRoot}/src", "args": [ - "--config=${workspaceRoot}/themes/jandedobbeleer.omp.json", + "--config=${workspaceRoot}/themes/cinnamon.omp.json", "--shell=pwsh", "--terminal-width=200", ] diff --git a/src/environment/shell.go b/src/environment/shell.go index df169edb..ea2a8657 100644 --- a/src/environment/shell.go +++ b/src/environment/shell.go @@ -164,7 +164,7 @@ type Environment interface { ExecutionTime() float64 Args() *Args BatteryInfo() ([]*battery.Battery, error) - WindowTitle(imageName, windowTitleRegex string) (string, error) + QueryWindowTitles(processName, windowTitleRegex string) (string, error) WindowsRegistryKeyValue(path string) (*WindowsRegistryValue, error) HTTPRequest(url string, timeout int, requestModifiers ...HTTPRequestModifier) ([]byte, error) IsWsl() bool diff --git a/src/environment/unix.go b/src/environment/unix.go index 448997c6..30cdf4b5 100644 --- a/src/environment/unix.go +++ b/src/environment/unix.go @@ -21,7 +21,7 @@ func (env *ShellEnvironment) Home() string { return os.Getenv("HOME") } -func (env *ShellEnvironment) WindowTitle(imageName, windowTitleRegex string) (string, error) { +func (env *ShellEnvironment) QueryWindowTitles(processName, windowTitleRegex string) (string, error) { return "", errors.New("not implemented") } diff --git a/src/environment/windows.go b/src/environment/windows.go index a221f221..c5807dfe 100644 --- a/src/environment/windows.go +++ b/src/environment/windows.go @@ -70,9 +70,9 @@ func (env *ShellEnvironment) Home() string { return home } -func (env *ShellEnvironment) WindowTitle(imageName, windowTitleRegex string) (string, error) { - defer env.trace(time.Now(), "WindowTitle", imageName, windowTitleRegex) - return WindowTitle(imageName, windowTitleRegex) +func (env *ShellEnvironment) QueryWindowTitles(processName, windowTitleRegex string) (string, error) { + defer env.trace(time.Now(), "WindowTitle", windowTitleRegex) + return queryWindowTitles(processName, windowTitleRegex) } func (env *ShellEnvironment) IsWsl() bool { diff --git a/src/environment/windows_win32.go b/src/environment/windows_win32.go index b4fb121d..b2944032 100644 --- a/src/environment/windows_win32.go +++ b/src/environment/windows_win32.go @@ -4,7 +4,6 @@ package environment import ( "errors" - "fmt" "oh-my-posh/regex" "strings" "syscall" @@ -14,99 +13,6 @@ import ( "golang.org/x/sys/windows" ) -// WindowsProcess is an implementation of Process for Windows. -type WindowsProcess struct { - pid int - ppid int - exe string -} - -// getImagePid returns the -func getImagePid(imageName string) ([]int, error) { - processes, err := processes() - if err != nil { - return nil, err - } - var pids []int - for i := 0; i < len(processes); i++ { - if strings.ToLower(processes[i].exe) == imageName { - pids = append(pids, processes[i].pid) - } - } - return pids, nil -} - -// WindowTitle returns the title of a window linked to a process name -func WindowTitle(imageName, windowTitleRegex string) (string, error) { - processPid, err := getImagePid(imageName) - if err != nil { - return "", nil - } - - // is a spotify process running? - // no: returns an empty string - if len(processPid) == 0 { - return "", nil - } - - // returns the first window of the first pid - _, windowTitle := GetWindowTitle(processPid[0], windowTitleRegex) - - return windowTitle, nil -} - -func newWindowsProcess(e *windows.ProcessEntry32) *WindowsProcess { - // Find when the string ends for decoding - end := 0 - for { - if e.ExeFile[end] == 0 { - break - } - end++ - } - - return &WindowsProcess{ - pid: int(e.ProcessID), - ppid: int(e.ParentProcessID), - exe: syscall.UTF16ToString(e.ExeFile[:end]), - } -} - -// Processes returns a snapshot of all the processes -// Taken and adapted from https://github.com/mitchellh/go-ps -func processes() ([]WindowsProcess, error) { - // get process table snapshot - handle, err := windows.CreateToolhelp32Snapshot(windows.TH32CS_SNAPPROCESS, 0) - if err != nil { - return nil, syscall.GetLastError() - } - defer func() { - _ = windows.CloseHandle(handle) - }() - - // get process infor by looping through the snapshot - var entry windows.ProcessEntry32 - entry.Size = uint32(unsafe.Sizeof(entry)) - err = windows.Process32First(handle, &entry) - if err != nil { - return nil, fmt.Errorf("error retrieving process info") - } - - results := make([]WindowsProcess, 0, 50) - for { - results = append(results, *newWindowsProcess(&entry)) - err := windows.Process32Next(handle, &entry) - if err != nil { - if err == syscall.ERROR_NO_MORE_FILES { - break - } - return nil, fmt.Errorf("Fail to syscall Process32Next: %v", err) - } - } - - return results, nil -} - // win32 specific code // win32 dll load and function definitions @@ -115,11 +21,14 @@ var ( procEnumWindows = user32.NewProc("EnumWindows") procGetWindowTextW = user32.NewProc("GetWindowTextW") procGetWindowThreadProcessID = user32.NewProc("GetWindowThreadProcessId") + + psapi = syscall.NewLazyDLL("psapi.dll") + getModuleBaseNameA = psapi.NewProc("GetModuleBaseNameA") ) -// EnumWindows call EnumWindows from user32 and returns all active windows +// enumWindows call enumWindows from user32 and returns all active windows // https://docs.microsoft.com/en-us/windows/win32/api/winuser/nf-winuser-enumwindows -func EnumWindows(enumFunc, lparam uintptr) (err error) { +func enumWindows(enumFunc, lparam uintptr) (err error) { r1, _, e1 := syscall.Syscall(procEnumWindows.Addr(), 2, enumFunc, lparam, 0) if r1 == 0 { if e1 != 0 { @@ -131,9 +40,9 @@ func EnumWindows(enumFunc, lparam uintptr) (err error) { return } -// GetWindowText returns the title and text of a window from a window handle +// getWindowText returns the title and text of a window from a window handle // https://docs.microsoft.com/en-us/windows/win32/api/winuser/nf-winuser-getwindowtextw -func GetWindowText(hwnd syscall.Handle, str *uint16, maxCount int32) (length int32, err error) { +func getWindowText(hwnd syscall.Handle, str *uint16, maxCount int32) (length int32, err error) { r0, _, e1 := syscall.Syscall(procGetWindowTextW.Addr(), 3, uintptr(hwnd), uintptr(unsafe.Pointer(str)), uintptr(maxCount)) length = int32(r0) if length == 0 { @@ -146,41 +55,64 @@ func GetWindowText(hwnd syscall.Handle, str *uint16, maxCount int32) (length int return } +func getWindowFileName(handle syscall.Handle) (string, error) { + var pid int + // get pid + _, _, err := procGetWindowThreadProcessID.Call(uintptr(handle), uintptr(unsafe.Pointer(&pid))) + if err != nil && err.Error() != "The operation completed successfully." { + return "", errors.New("unable to get window process pid") + } + const query = windows.PROCESS_QUERY_INFORMATION | windows.PROCESS_VM_READ + h, err := windows.OpenProcess(query, false, uint32(pid)) + if err != nil { + return "", errors.New("unable to open window process") + } + buf := [1024]byte{} + length, _, err := getModuleBaseNameA.Call(uintptr(h), 0, uintptr(unsafe.Pointer(&buf)), 1024) + if err != nil && err.Error() != "The operation completed successfully." { + return "", errors.New("unable to get window process name") + } + filename := string(buf[:length]) + return strings.ToLower(filename), nil +} + // GetWindowTitle searches for a window attached to the pid -func GetWindowTitle(pid int, windowTitleRegex string) (syscall.Handle, string) { - var hwnd syscall.Handle +func queryWindowTitles(processName, windowTitleRegex string) (string, error) { var title string - - // callback fro EnumWindows - cb := syscall.NewCallback(func(h syscall.Handle, p uintptr) uintptr { - var prcsID int - // get pid - _, _, _ = procGetWindowThreadProcessID.Call(uintptr(h), uintptr(unsafe.Pointer(&prcsID))) - // check if pid matches spotify pid - if prcsID == pid { - b := make([]uint16, 200) - _, err := GetWindowText(h, &b[0], int32(len(b))) - if err != nil { - // ignore the error - return 1 // continue enumeration - } - title = syscall.UTF16ToString(b) - if regex.MatchString(windowTitleRegex, title) { - // will cause EnumWindows to return 0 (error) - // but we don't want to enumerate all windows since we got what we want - hwnd = h - return 0 - } + // callback for EnumWindows + cb := syscall.NewCallback(func(handle syscall.Handle, pointer uintptr) uintptr { + fileName, err := getWindowFileName(handle) + if err != nil { + // ignore the error and continue enumeration + return 1 + } + if processName != fileName { + // ignore the error and continue enumeration + return 1 + } + b := make([]uint16, 200) + _, err = getWindowText(handle, &b[0], int32(len(b))) + if err != nil { + // ignore the error and continue enumeration + return 1 + } + title = syscall.UTF16ToString(b) + if regex.MatchString(windowTitleRegex, title) { + // will cause EnumWindows to return 0 (error) + // but we don't want to enumerate all windows since we got what we want + return 0 } - return 1 // continue enumeration }) // Enumerates all top-level windows on the screen // The error is not checked because if EnumWindows is stopped bofere enumerating all windows // it returns 0(error occurred) instead of 1(success) // In our case, title will equal "" or the title of the window anyway - _ = EnumWindows(cb, 0) - return hwnd, title + _ = enumWindows(cb, 0) + if len(title) == 0 { + return "", errors.New("no matching window title found") + } + return title, nil } // Return the windows handles corresponding to the names of the root registry keys. diff --git a/src/mock/environment.go b/src/mock/environment.go index 6f03cbf4..c9eb270f 100644 --- a/src/mock/environment.go +++ b/src/mock/environment.go @@ -121,8 +121,8 @@ func (env *MockedEnvironment) Shell() string { return args.String(0) } -func (env *MockedEnvironment) WindowTitle(imageName, windowTitleRegex string) (string, error) { - args := env.Called(imageName) +func (env *MockedEnvironment) QueryWindowTitles(processName, windowTitleRegex string) (string, error) { + args := env.Called(processName, windowTitleRegex) return args.String(0), args.Error(1) } diff --git a/src/segments/spotify_windows.go b/src/segments/spotify_windows.go index 6731e8ca..cddb605d 100644 --- a/src/segments/spotify_windows.go +++ b/src/segments/spotify_windows.go @@ -9,20 +9,32 @@ import ( func (s *Spotify) Enabled() bool { // search for spotify window to retrieve the title // Can be either "Spotify xxx" or the song name "Candlemass - Spellbreaker" - spotifyWindowTitle, err := s.env.WindowTitle("spotify.exe", "^(Spotify.*)|(.*\\s-\\s.*)$") + windowTitle, err := s.env.QueryWindowTitles("spotify.exe", "^(Spotify.*)|(.*\\s-\\s.*)$") + if err == nil { + return s.parseSpotifyTitle(windowTitle, " - ") + } + windowTitle, err = s.env.QueryWindowTitles("msedge.exe", "^(Spotify.*)") if err != nil { return false } + return s.parseWebSpotifyTitle(windowTitle) +} - if !strings.Contains(spotifyWindowTitle, " - ") { +func (s *Spotify) parseWebSpotifyTitle(windowTitle string) bool { + windowTitle = strings.TrimPrefix(windowTitle, "Spotify - ") + return s.parseSpotifyTitle(windowTitle, " · ") +} + +func (s *Spotify) parseSpotifyTitle(windowTitle, separator string) bool { + if !strings.Contains(windowTitle, separator) { s.Status = stopped return false } - infos := strings.Split(spotifyWindowTitle, " - ") + infos := strings.Split(windowTitle, separator) s.Artist = infos[0] // remove first element and concat others(a song can contains also a " - ") - s.Track = strings.Join(infos[1:], " - ") + s.Track = strings.Join(infos[1:], separator) s.Status = playing s.resolveIcon() return true diff --git a/src/segments/spotify_windows_test.go b/src/segments/spotify_windows_test.go index 38f721e9..6ea59852 100644 --- a/src/segments/spotify_windows_test.go +++ b/src/segments/spotify_windows_test.go @@ -11,47 +11,72 @@ import ( "github.com/stretchr/testify/assert" ) -type spotifyArgs struct { - title string - runError error +func TestSpotifyWindowsNative(t *testing.T) { + cases := []struct { + Case string + ExpectedString string + ExpectedEnabled bool + Title string + Error error + }{ + { + Case: "Playing", + ExpectedString: "\ue602 Candlemass - Spellbreaker", + ExpectedEnabled: true, + Title: "Candlemass - Spellbreaker", + }, + { + Case: "Stopped", + ExpectedEnabled: false, + Title: "Spotify premium", + }, + } + for _, tc := range cases { + env := new(mock.MockedEnvironment) + env.On("QueryWindowTitles", "spotify.exe", "^(Spotify.*)|(.*\\s-\\s.*)$").Return(tc.Title, tc.Error) + env.On("QueryWindowTitles", "msedge.exe", "^(Spotify.*)").Return("", errors.New("not implemented")) + s := &Spotify{ + env: env, + props: properties.Map{}, + } + assert.Equal(t, tc.ExpectedEnabled, s.Enabled()) + if tc.ExpectedEnabled { + assert.Equal(t, tc.ExpectedString, renderTemplate(env, s.Template(), s)) + } + } } -func bootStrapSpotifyWindowsTest(args *spotifyArgs) *Spotify { - env := new(mock.MockedEnvironment) - env.On("WindowTitle", "spotify.exe").Return(args.title, args.runError) - s := &Spotify{ - env: env, - props: properties.Map{}, +func TestSpotifyWindowsPWA(t *testing.T) { + cases := []struct { + Case string + ExpectedString string + ExpectedEnabled bool + Title string + Error error + }{ + { + Case: "Playing", + ExpectedString: "\ue602 Snow in Stockholm - Sarah, the Illstrumentalist", + ExpectedEnabled: true, + Title: "Spotify - Snow in Stockholm · Sarah, the Illstrumentalist", + }, + { + Case: "Stopped", + ExpectedEnabled: false, + Title: "Spotify - Web Player", + }, + } + for _, tc := range cases { + env := new(mock.MockedEnvironment) + env.On("QueryWindowTitles", "spotify.exe", "^(Spotify.*)|(.*\\s-\\s.*)$").Return("", errors.New("not implemented")) + env.On("QueryWindowTitles", "msedge.exe", "^(Spotify.*)").Return(tc.Title, tc.Error) + s := &Spotify{ + env: env, + props: properties.Map{}, + } + assert.Equal(t, tc.ExpectedEnabled, s.Enabled()) + if tc.ExpectedEnabled { + assert.Equal(t, tc.ExpectedString, renderTemplate(env, s.Template(), s)) + } } - return s -} - -func TestSpotifyWindowsEnabledAndSpotifyNotRunning(t *testing.T) { - args := &spotifyArgs{ - runError: errors.New(""), - } - s := bootStrapSpotifyWindowsTest(args) - assert.Equal(t, false, s.Enabled()) -} - -func TestSpotifyWindowsEnabledAndSpotifyPlaying(t *testing.T) { - args := &spotifyArgs{ - title: "Candlemass - Spellbreaker", - } - env := new(mock.MockedEnvironment) - env.On("WindowTitle", "spotify.exe").Return(args.title, args.runError) - s := &Spotify{ - env: env, - props: properties.Map{}, - } - assert.Equal(t, true, s.Enabled()) - assert.Equal(t, "\ue602 Candlemass - Spellbreaker", renderTemplate(env, s.Template(), s)) -} - -func TestSpotifyWindowsEnabledAndSpotifyStopped(t *testing.T) { - args := &spotifyArgs{ - title: "Spotify premium", - } - s := bootStrapSpotifyWindowsTest(args) - assert.Equal(t, false, s.Enabled()) } diff --git a/src/segments/spotify_wsl_test.go b/src/segments/spotify_wsl_test.go index 388370b5..a8ccbeee 100644 --- a/src/segments/spotify_wsl_test.go +++ b/src/segments/spotify_wsl_test.go @@ -51,7 +51,8 @@ func TestSpotifyWsl(t *testing.T) { Case: "tasklist.exe not in path", ExpectedString: "-", ExpectedEnabled: false, - ExecOutput: ""}, + ExecOutput: "", + }, } for _, tc := range cases { env := new(mock.MockedEnvironment)