feat(windows): registry query segment

docs/docs/segment-winreg.md

@@ -0,0 +1,46 @@
+---
+id: winreg
+title: Windows Registry Key Query
+sidebar_label: Windows Registry Key Query
+---
+
+## What
+
+Display the content of the requested Windows registry key.
+
+Supported registry key types:
+
+- String
+- DWORD (displayed in upper-case 0x hex)
+
+## Sample Configuration
+
+```json
+  {
+    "type": "winreg",
+    "style": "powerline",
+    "powerline_symbol": "\uE0B0",
+    "foreground": "#ffffff",
+    "background": "#444444",
+    "properties": {
+      "path": "HKLM\\software\\microsoft\\windows nt\\currentversion",
+      "key":"buildlab",
+      "template":"{{ if .Value }}{{ .Value }}{{ else }}unknown{{ end }}",
+      "prefix": " \uE62A "
+    }
+  },
+```
+
+## Properties
+
+- path: `string` - registry path to the desired key using backslashes and with a valid root HKEY name.
+- key: `string` - the key to read from the `path location.  If `""`, will read the default value.
+- template: `string` - a go [text/template][go-text-template] template extended
+  with [sprig][sprig] utilizing the properties below.
+
+## Template Properties
+
+- .Value: `string` - The result of your query
+
+[go-text-template]: https://golang.org/pkg/text/template/
+[sprig]: https://masterminds.github.io/sprig/

docs/sidebars.js

@@ -73,6 +73,7 @@ module.exports = {
         "text",
         "time",
         "wifi",
+        "winreg",
         "ytm",
       ],
     },

src/environment.go

@@ -80,6 +80,7 @@ type environmentInfo interface {
 	getBatteryInfo() ([]*battery.Battery, error)
 	getShellName() string
 	getWindowTitle(imageName, windowTitleRegex string) (string, error)
+	getWindowsRegistryKeyValue(regPath, regKey string) (string, error)
 	doGet(url string, timeout int) ([]byte, error)
 	hasParentFilePath(path string) (fileInfo *fileInfo, err error)
 	isWsl() bool

src/environment_unix.go

@@ -59,3 +59,7 @@ func (env *environment) getCachePath() string {
 	}
 	return env.homeDir()
 }
+
+func (env *environment) getWindowsRegistryKeyValue(regPath, regKey string) (string, error) {
+	return "", errors.New("not implemented")
+}

src/environment_windows.go

@@ -3,9 +3,13 @@
 package main
 
 import (
+	"errors"
+	"fmt"
 	"os"
+	"strings"
 	"syscall"
 	"time"
+	"unsafe"
 
 	"github.com/Azure/go-ansiterm/winterm"
 	"golang.org/x/sys/windows"
@@ -99,3 +103,98 @@ func (env *environment) getCachePath() string {
 	}
 	return env.homeDir()
 }
+
+//
+// Takes a registry path like "HKLM\Software\Microsoft\Windows NT\CurrentVersion" and a key under that path like "CurrentVersion" (or "" if the (Default) key is required).
+// Returns a bool and string:
+//
+//   true and the retrieved value formatted into a string if successful.
+//   false and the string will be the error
+//
+func (env *environment) getWindowsRegistryKeyValue(regPath, regKey string) (string, error) {
+	env.trace(time.Now(), "getWindowsRegistryKeyValue", regPath, regKey)
+
+	// Extract root HK value and turn it into a windows.Handle to open the key.
+	regPathParts := strings.SplitN(regPath, "\\", 2)
+
+	regRootHKeyHandle := getHKEYHandleFromAbbrString(regPathParts[0])
+	if regRootHKeyHandle == 0 {
+		errorLogMsg := fmt.Sprintf("Error, Supplied root HKEY value not valid: '%s'", regPathParts[0])
+		env.log(Error, "getWindowsRegistryKeyValue", errorLogMsg)
+		return "", errors.New(errorLogMsg)
+	}
+
+	// Second part of split is registry path after HK part - needs to be UTF16 to pass to the windows. API
+	regPathUTF16, regPathUTF16ConversionErr := windows.UTF16FromString(regPathParts[1])
+	if regPathUTF16ConversionErr != nil {
+		errorLogMsg := fmt.Sprintf("Error, Could not convert supplied path '%s' to UTF16, error: '%s'", regPathParts[1], regPathUTF16ConversionErr)
+		env.log(Error, "getWindowsRegistryKeyValue", errorLogMsg)
+		return "", errors.New(errorLogMsg)
+	}
+
+	// Ok - open it..
+	var hKeyHandle windows.Handle
+	regOpenErr := windows.RegOpenKeyEx(regRootHKeyHandle, &regPathUTF16[0], 0, windows.KEY_READ, &hKeyHandle)
+	if regOpenErr != nil {
+		errorLogMsg := fmt.Sprintf("Error RegOpenKeyEx opening registry path to '%s', error: '%s'", regPath, regOpenErr)
+		env.log(Error, "getWindowsRegistryKeyValue", errorLogMsg)
+		return "", errors.New(errorLogMsg)
+	}
+	// Success - from here on out, when returning make sure to close that reg key with a deferred call to close:
+	defer func() {
+		err := windows.RegCloseKey(hKeyHandle)
+		if err != nil {
+			env.log(Error, "getWindowsRegistryKeyValue", fmt.Sprintf("Error closing registry key: %s", err))
+		}
+	}()
+
+	// Again - need UTF16 of the key for the API:
+	regKeyUTF16, regKeyUTF16ConversionErr := windows.UTF16FromString(regKey)
+	if regKeyUTF16ConversionErr != nil {
+		errorLogMsg := fmt.Sprintf("Error, could not convert supplied key '%s' to UTF16, error: '%s'", regKey, regKeyUTF16ConversionErr)
+		env.log(Error, "getWindowsRegistryKeyValue", errorLogMsg)
+		return "", errors.New(errorLogMsg)
+	}
+
+	// Two stage way to get the key value - query once to get size - then allocate and query again to fill it.
+	var keyBufType uint32
+	var keyBufSize uint32
+
+	regQueryErr := windows.RegQueryValueEx(hKeyHandle, &regKeyUTF16[0], nil, &keyBufType, nil, &keyBufSize)
+	if regQueryErr != nil {
+		errorLogMsg := fmt.Sprintf("Error calling RegQueryValueEx to retrieve key data size with error '%s'", regQueryErr)
+		env.log(Error, "getWindowsRegistryKeyValue", errorLogMsg)
+		return "", errors.New(errorLogMsg)
+	}
+
+	// Alloc and fill...
+	var keyBuf = make([]byte, keyBufSize)
+
+	regQueryErr = windows.RegQueryValueEx(hKeyHandle, &regKeyUTF16[0], nil, &keyBufType, &keyBuf[0], &keyBufSize)
+	if regQueryErr != nil {
+		errorLogMsg := fmt.Sprintf("Error calling RegQueryValueEx to retrieve key data with error '%s'", regQueryErr)
+		env.log(Error, "getWindowsRegistryKeyValue", errorLogMsg)
+		return "", errors.New(errorLogMsg)
+	}
+
+	// Format result into a string, depending on type.  (future refactor - move this out into it's own function)
+	switch keyBufType {
+	case windows.REG_SZ:
+		var uint16p *uint16
+		uint16p = (*uint16)(unsafe.Pointer(&keyBuf[0])) // nasty casty
+
+		valueString := windows.UTF16PtrToString(uint16p)
+		env.log(Debug, "getWindowsRegistryKeyValue", fmt.Sprintf("success, string: %s", valueString))
+		return valueString, nil
+	case windows.REG_DWORD:
+		var uint32p *uint32
+		uint32p = (*uint32)(unsafe.Pointer(&keyBuf[0])) // more casting goodness
+
+		valueString := fmt.Sprintf("0x%08X", *uint32p)
+		env.log(Debug, "getWindowsRegistryKeyValue", fmt.Sprintf("success, DWORD, formatted as string: %s", valueString))
+		return valueString, nil
+	default:
+		errorLogMsg := fmt.Sprintf("Error, no formatter for REG_? type:%d, data size:%d bytes", keyBufType, keyBufSize)
+		return "", errors.New(errorLogMsg)
+	}
+}

src/environment_windows_win32.go

@@ -179,3 +179,22 @@ func GetWindowTitle(pid int, windowTitleRegex string) (syscall.Handle, string) {
 	_ = EnumWindows(cb, 0)
 	return hwnd, title
 }
+
+// Return the windows handles corresponding to the names of the root registry keys.
+// A returned value of 0 means there was no match.
+func getHKEYHandleFromAbbrString(abbr string) windows.Handle {
+	switch abbr {
+	case "HKCR", "HKEY_CLASSES_ROOT":
+		return windows.HKEY_CLASSES_ROOT
+	case "HKCC", "HKEY_CURRENT_CONFIG":
+		return windows.HKEY_CURRENT_CONFIG
+	case "HKCU", "HKEY_CURRENT_USER":
+		return windows.HKEY_CURRENT_USER
+	case "HKLM", "HKEY_LOCAL_MACHINE":
+		return windows.HKEY_LOCAL_MACHINE
+	case "HKU", "HKEY_USERS":
+		return windows.HKEY_USERS
+	}
+
+	return 0
+}

src/segment.go

@@ -134,6 +134,8 @@ const (
 	Nightscout SegmentType = "nightscout"
 	// WiFi writes details about the current WiFi connection
 	WiFi SegmentType = "wifi"
+	// WinReg queries the Windows registry.
+	WinReg SegmentType = "winreg"
 )
 
 func (segment *Segment) string() string {
@@ -269,6 +271,7 @@ func (segment *Segment) mapSegmentWithWriter(env environmentInfo) error {
 		PHP:           &php{},
 		Nightscout:    &nightscout{},
 		WiFi:          &wifi{},
+		WinReg:        &winreg{},
 	}
 	if writer, ok := functions[segment.Type]; ok {
 		props := &properties{

src/segment_path_test.go

@@ -129,6 +129,11 @@ func (env *MockedEnvironment) getWindowTitle(imageName, windowTitleRegex string)
 	return args.String(0), args.Error(1)
 }
 
+func (env *MockedEnvironment) getWindowsRegistryKeyValue(regPath, regKey string) (string, error) {
+	args := env.Called(regPath, regKey)
+	return args.String(0), args.Error(1)
+}
+
 func (env *MockedEnvironment) doGet(url string, timeout int) ([]byte, error) {
 	args := env.Called(url)
 	return args.Get(0).([]byte), args.Error(1)

src/segment_winreg.go

@@ -0,0 +1,51 @@
+package main
+
+type winreg struct {
+	props *properties
+	env   environmentInfo
+
+	Value string
+}
+
+const (
+	// path from the supplied root under which the key exists
+	RegistryPath Property = "path"
+	// key within full reg path formed from two above
+	RegistryKey Property = "key"
+)
+
+func (wr *winreg) init(props *properties, env environmentInfo) {
+	wr.props = props
+	wr.env = env
+}
+
+func (wr *winreg) enabled() bool {
+	if wr.env.getRuntimeGOOS() != windowsPlatform {
+		return false
+	}
+
+	registryPath := wr.props.getString(RegistryPath, "")
+	registryKey := wr.props.getString(RegistryKey, "")
+
+	var err error
+	wr.Value, err = wr.env.getWindowsRegistryKeyValue(registryPath, registryKey)
+	return err == nil
+}
+
+func (wr *winreg) string() string {
+	segmentTemplate := wr.props.getString(SegmentTemplate, "{{ .Value }}")
+	return wr.templateString(segmentTemplate)
+}
+
+func (wr *winreg) templateString(segmentTemplate string) string {
+	template := &textTemplate{
+		Template: segmentTemplate,
+		Context:  wr,
+		Env:      wr.env,
+	}
+	text, err := template.render()
+	if err != nil {
+		return err.Error()
+	}
+	return text
+}

src/segment_winreg_test.go

@@ -0,0 +1,52 @@
+package main
+
+import (
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRegQueryEnabled(t *testing.T) {
+	cases := []struct {
+		CaseDescription string
+		Path            string
+		Key             string
+		ExpectedSuccess bool
+		Output          string
+		Err             error
+	}{
+		{
+			CaseDescription: "Error",
+			Path:            "HKLLM\\Software\\Microsoft\\Windows NT\\CurrentVersion",
+			Key:             "ProductName",
+			ExpectedSuccess: false,
+			Err:             errors.New("No match"),
+		},
+		{
+			CaseDescription: "Value",
+			Path:            "HKLM\\Software\\Microsoft\\Windows NT\\CurrentVersion",
+			Key:             "InstallTime",
+			ExpectedSuccess: true,
+			Output:          "no formatter",
+		},
+	}
+
+	for _, tc := range cases {
+		env := new(MockedEnvironment)
+		env.On("getRuntimeGOOS", nil).Return(windowsPlatform)
+		env.On("getWindowsRegistryKeyValue", tc.Path, tc.Key).Return(tc.Output, tc.Err)
+		props := &properties{
+			values: map[Property]interface{}{
+				RegistryPath: tc.Path,
+				RegistryKey:  tc.Key,
+			},
+		}
+		r := &winreg{
+			env:   env,
+			props: props,
+		}
+
+		assert.Equal(t, r.enabled(), tc.ExpectedSuccess, tc.CaseDescription)
+	}
+}

themes/schema.json

@@ -181,7 +181,8 @@
             "sysinfo",
             "angular",
             "php",
-            "wifi"
+            "wifi",
+            "regquery"
           ]
         },
         "style": {
@@ -1718,6 +1719,38 @@
               }
             }
           }
+        },
+        {
+          "if": {
+            "properties": {
+              "type": { "const": "winreg" }
+            }
+          },
+          "then": {
+            "title": "Windows Registry Query",
+            "description": "https://ohmyposh.dev/docs/winreg",
+            "properties": {
+              "properties": {
+                "properties": {
+                  "template": {
+                    "$ref": "#/definitions/template"
+                  },
+                  "path": {
+                    "type": "string",
+                    "title": "Registry Path",
+                    "description": "The path under which the registy key lives (case insensitive, must use backslashes), e.g. HKLM\\Software\\Microsoft\\Windows NT\\CurrentVersion",
+                    "default": ""
+                  },
+                  "key": {
+                    "type": "string",
+                    "title": "Registry Key",
+                    "description": "The key under he registry path to get (case insensitive).  If left blank, will get the value of the (Default) key in the registry_path",
+                    "default": ""
+                  }
+                }
+              }
+            }
+          }
         }
       ]
     }