Update fleetContextMenu.py

fleetContextMenu.py

@@ -1,53 +1,72 @@
-
 import os
+import sys
 import winreg
+import ctypes
+import time
+
+def is_admin():
+    try:
+        return ctypes.windll.shell32.IsUserAnAdmin()
+    except:
+        return False
+
+def run_as_admin():
+    ctypes.windll.shell32.ShellExecuteW(
+        None, "runas", sys.executable, " ".join(sys.argv), None, 1
+    )
 
 # Define paths and menu text
 menu_text = "Open with Fleet"
-app_path = r"C:\Users\dower\AppData\Local\Programs\Fleet\Fleet.exe"
-icon_path = r"C:\Users\dower\AppData\Local\Programs\Fleet\Fleet.ico"
+appdata_local_path = os.path.join(os.getenv('LOCALAPPDATA'), 'Programs')
 
-def add_context_menu():
-    """Adds a new context menu entry for folders."""
+# Specify the rest of the path to your executable
+executable_path = 'Fleet\Fleet.exe'
 
-    # Open registry key
-    key = winreg.OpenKey(winreg.HKEY_CLASSES_ROOT, r"Directory\Background\shell", 0, winreg.KEY_WRITE)
+# Combine the paths
+app_path = os.path.join(appdata_local_path, executable_path)
 
-    # Create subkey for the menu entry
-    subkey = winreg.CreateKey(key, menu_text)
+def add_context_menu(location):
+    """Adds a new context menu entry for folders."""
 
-    # **Change:** Use f-string with curly braces for command
-    winreg.SetValueEx(subkey, "command", 0, winreg.REG_SZ, f"\"{app_path}\" --dir \"%1\"")
+    # Open registry key
+    key_path = f"Directory\\{location}\\Open with Fleet"
+    key = winreg.CreateKey(winreg.HKEY_CLASSES_ROOT, key_path)
+    subkey = winreg.CreateKey(key, "command")
 
-    # Set default icon for the menu entry
-    winreg.SetValueEx(subkey, "Icon", 0, winreg.REG_SZ, f"{app_path},0")
+    # Set default values
+    winreg.SetValueEx(key, 'Icon', 0, winreg.REG_SZ, f"\"{app_path}\",0")
+    winreg.SetValueEx(key, None, 0, winreg.REG_SZ, menu_text)
+    winreg.SetValueEx(subkey, None, 0, winreg.REG_SZ, f"\"{app_path}\" --dir \"%V\"")
 
     # Close registry keys
     winreg.CloseKey(subkey)
     winreg.CloseKey(key)
 
-def remove_context_menu():
+def remove_context_menu(location):
     """Removes the existing context menu entry."""
+    key_path = f"Directory\\{location}\\Open with Fleet"
 
     try:
-        # Open registry key
-        key = winreg.OpenKey(winreg.HKEY_CLASSES_ROOT, r"Directory\Background\shell", 0, winreg.KEY_DELETE)
+        # Delete the registry key and its subkeys
+        winreg.DeleteKey(winreg.HKEY_CLASSES_ROOT, key_path)
+    except FileNotFoundError:
+        print(f"Registry key not found: {key_path}")
+    except Exception as e:
+        print(f"Error removing context menu entry: {e}")
 
-        # Delete the subkey for the menu entry
-        winreg.DeleteKey(key, menu_text)
+# Check if running as admin
+if not is_admin():
+    # Re-run the script as admin
+    run_as_admin()
+    sys.exit()
+
+# Remove existing entries
+remove_context_menu("Background\\shell")
+remove_context_menu("shell")
 
-        # Close registry key
-        winreg.CloseKey(key)
-    except FileNotFoundError:
-        pass
-
-try:
-    # Check if existing entry needs to be removed
-    if menu_text in os.listdir("HKEY_CLASSES_ROOT\\Directory\\Background\\shell"):
-        remove_context_menu()
-except:
-    print("skipping remove")
 # Add new context menu entry
-add_context_menu()
+add_context_menu("Background\\shell")
+add_context_menu("shell")
 
 print(f"Successfully added '{menu_text}' context menu entry for folders.")
+time.sleep(3)