Windows ContextMenu Editing Through Custom Registry Module

reg.py

import winreg

class BinaryRegValue:
    def __init__(self, value):
        self.value = value

    def __repr__(self):
        return f"<Binary value={repr(self.value)[1:]}>"

class StringRegValue:
    def __init__(self, value):
        self.value = value

    def __repr__(self):
        return f"<String value='{self.value}'>"

class NumberRegValue:
    def __init__(self, value, bits):
        self.value = value
        self.bits = bits

    def __repr__(self):
        return f"<Number{self.bits} value={self.value}>"

class StringArrayRegValue:
    def __init__(self, value):
        self.value = value

    def __repr__(self):
        return f"<StringArray value={self.value}>"

class NonExistantRegKey:
    def __init__(self, container, key):
        self.container = container  # Reference to RegKey
        self.key = key              # The key to populate

    def make(self):
        new_key = RegKey.open(self.key, True, self.container)
        self.container.data[self.key] = new_key
        return new_key

class RegKeyObject:
    def __init__(self, handle):
        self.handle = handle
    
    def __getitem__(self, key):
        if key == "":
            return winreg.QueryValue(self.handle, None)
        else:
            try:
                return winreg.QueryValueEx(self.handle, key)[0]
            except OSError:
                return None

    def __setitem__(self, key, regobject):
        typevalue = type(regobject)
        if typevalue == StringRegValue:
                winreg.SetValueEx(self.handle, key, 0, 1, regobject.value)
        elif typevalue == BinaryRegValue:
                winreg.SetValueEx(self.handle, key, 0, 3, regobject.value)

class RegKey:
    @classmethod
    def open(cls, subkey, create=False, parentkey=winreg.HKEY_CURRENT_USER):
        if type(parentkey) == RegKey:
            parent = parentkey.parentkey
            subkey = cls.join_path(parentkey.keyname, subkey)
        else:
            parent = parentkey
        try:
            handle = winreg.OpenKey(parent, subkey, 0, winreg.KEY_WRITE | winreg.KEY_READ)
        except OSError:
            if create:
                handle = winreg.CreateKey(parent, subkey)
            else:
                return None
        return cls(parent, handle, subkey)

    @classmethod
    def join_path(cls, path1, path2):
        if path1 == None:
            return path2
        else:
            return path1 + "\\" + path2

    def create(self, subkey):
        handle = winreg.CreateKey(self.handle, subkey)
        subkey = self.keyname + "\\" + subkey
        return RegKey(self.parentkey, handle, subkey)

    def __init__(self, parentkey, handle, keyname):
        super().__setattr__('data', {})  # Avoid triggering __setattr__
        self.field = RegKeyObject(handle)
        self.parentkey = parentkey
        self.handle = handle
        self.keyname = keyname

    def subkeys(self):
        subkeys = []
        try:
            children, _, _ = winreg.QueryInfoKey(self.handle)
            for index in range(children):
                subkeys.append(winreg.EnumKey(self.handle, index))
        except FileNotFoundError:
            pass
        return subkeys

    def enum_values(self):
        values = {}
        i = 0
        while True:
            try:
                name, value, vtype = winreg.EnumValue(self.handle, i)
                match vtype:
                    case 1:
                        values[name] = StringRegValue(value)
                    case 3:
                        values[name] = BinaryRegValue(value)
                    case 4: # 32-bit word
                        values[name] = NumberRegValue(value, 32)
                    case 11: # 64-bit word
                        values[name] = NumberRegValue(value, 64)
                    case 7:
                        values[name] = StringArrayRegValue(value)
                    case _:
                        print(vtype)
            except OSError:
                return values
            i += 1

    def enum_keys(self):
        keys = []
        i = 0
        while True:
            try:
                childkeyname = winreg.EnumKey(self.handle, i)
                childkey = self.open(childkeyname, parentkey=self)
                children, _, _ = winreg.QueryInfoKey(childkey.handle)
                keys.append(childkey.keyname)
                if children > 0:
                    keys.extend(childkey.enum_keys())
                i += 1
            except WindowsError:
                return keys

    def children(self):
        keys = []
        i = 0
        while True:
            try:
                childkeyname = winreg.EnumKey(self.handle, i)
                if self.keyname != None:
                    keys.append(self.keyname + "\\" + childkeyname)
                else:
                    keys.append(childkeyname)
                i += 1
            except WindowsError:
                return keys
            

    def delete(self):
        subkeys = self.enum_keys()
        keys = [self.keyname]
        keys.extend(subkeys)

        for i in range(len(keys)-1, -1, -1):
            print("Deleting:", keys[i])
            winreg.DeleteKey(self.parentkey, keys[i])

    def __setattr__(self, key, value):
        if key == 'data':
            super().__setattr__(key, value)  # Set internal attribute
        else:
            self.data[key] = value  # Store in the internal dict

    def __getattr__(self, key):
        if key not in self.data:
            regkey = self.open(key, parentkey=self)
            if regkey != None:
                self.data[key] = regkey
            else:
                # Return a NonExistantRegKey that can populate itself when needed
                placeholder = NonExistantRegKey(self, key)
                self.data[key] = placeholder
        return self.data[key]

UserReg = RegKey.open(None)

if __name__ == "__main__":
    dir = UserReg.Software.Classes.Directory
    gnu = RegKey.open("GNU", True, dir.shell)
    command = RegKey.open("GNU\\command", True, dir.shell)
    gnu.field["icon"] = StringRegValue(r"C:\Users\kaecy\Data\Icons\road-block.ico")
    command.field[""] = StringRegValue("wsl --cd \"%v\"")

    for reg in dir.shell.children():
        print(reg)

UserContextMenu.py

import sys
import reg

def get_config():
  dictdata = {}
  dir = reg.RegKey.open(r"Software\Classes\Directory")
  shell = reg.RegKey.open("shell", True, dir)
  for subkey in shell.children():
    ok = reg.RegKey.open(subkey, parentkey=shell)
    dictdata[subkey] = {}
    cmd = ok.command
    if type(cmd) != reg.NonExistantRegKey:
      dictdata[subkey]["command"] = cmd.field[""]
    dictdata[subkey]["icon"] = ok.field["icon"]
  return dictdata

class UserContextMenu:
  def __init__(self):
    self.config = get_config()

  def delete(self, name):
    rk = reg.open_key(winreg.HKEY_CURRENT_USER, rf"Software\Classes\Directory\shell\{name}")
    if rk:
      rk.delete()
    else:
      print("Delete: context action name not found.")
    if name in self.config:
      del self.config[name]

  def write(self, name, command, icon=None):
    ok = winreg.CreateKey(winreg.HKEY_CURRENT_USER, rf"Software\Classes\Directory\shell\{name}")
    winreg.SetValue(ok, "command", 1, command)
    self.config[name] = {}
    self.config[name]["command"] = command
    if icon:
      winreg.SetValue(ok, icon)
      self.config[name]["icon"] = icon

  def print(self):
    for key, values in self.config.items():
      print(key, end=":\n")
      for key, value in values.items():
        print(" " * (10 - len(key)), key + ":", value)

menu = UserContextMenu()
menu.print()