Merge pull request mandiant#1221 from mandiant/fix-internet-detect

Fix Internet Detector Background/Taskbar Issues

packages/internet_detector.vm/internet_detector.vm.nuspec

@@ -2,7 +2,7 @@
 <package xmlns="http://schemas.microsoft.com/packaging/2015/06/nuspec.xsd">
   <metadata>
     <id>internet_detector.vm</id>
-    <version>1.0.0.20241216</version>
+    <version>1.0.0.20241217</version>
     <authors>Elliot Chernofsky and Ana Martinez Gomez</authors>
     <description>Tool that changes the background and a taskbar icon if it detects internet connectivity</description>
     <dependencies>

packages/internet_detector.vm/tools/chocolateyinstall.ps1

@@ -21,4 +21,9 @@ Start-Process -FilePath 'cmd.exe' -WorkingDirectory $toolDir -ArgumentList "/c p
 $imagesPath = Join-Path $packageToolDir "images"
 Copy-Item "$imagesPath\*" ${Env:VM_COMMON_DIR} -Force
 
-VM-Install-Shortcut -toolName $toolName -category $category -executablePath "$toolDir/$toolName.exe"
+VM-Install-Shortcut -toolName $toolName -category $category -executablePath "$toolDir\$toolName.exe"
+
+# Create scheduled task for tool to run every 2 minutes.
+$action = New-ScheduledTaskAction -Execute "$toolDir\$toolName.exe"
+$trigger = New-ScheduledTaskTrigger -Once -At (Get-Date) -RepetitionInterval (New-TimeSpan -Minutes 1)
+Register-ScheduledTask -Action $action -Trigger $trigger -TaskName 'Internet Detector' -Force

packages/internet_detector.vm/tools/internet_detector.pyw

@@ -28,7 +28,6 @@ import time
 import os
 import re
 
-
 # Define constants
 CHECK_INTERVAL = 2  # Seconds
 CONNECT_TEST_URL_AND_RESPONSES = {
@@ -48,6 +47,7 @@ ICON_INDICATOR_ON = os.path.join(os.environ.get("VM_COMMON_DIR"), "indicator_on.
 ICON_INDICATOR_OFF = os.path.join(os.environ.get("VM_COMMON_DIR"), "indicator_off.ico")
 DEFAULT_BACKGROUND = os.path.join(os.environ.get("VM_COMMON_DIR"), "background.png")
 INTERNET_BACKGROUND = os.path.join(os.environ.get("VM_COMMON_DIR"), "background-internet.png")
+REGISTRY_PATH = r"Software\InternetDetector"
 
 # Global variables
 tray_icon = None
@@ -62,6 +62,7 @@ default_palette = None
 # Win32 API icon handles
 hicon_indicator_off = None
 hicon_indicator_on = None
+mutex = None
 
 def is_already_running():
     global mutex
@@ -75,31 +76,35 @@ def is_already_running():
         return False  # Assume not running if error
 
 def signal_handler(sig, frame):
-    global check_thread, tray_icon_thread, tray_icon
+    global check_thread, tray_icon_thread, tray_icon, mutex
     print("Ctrl+C detected. Exiting...")
     stop_event.set()  # Signal the background thread to stop
     if check_thread:
         check_thread.join()
     if tray_icon_thread:
         tray_icon_thread.join()
     if tray_icon:
-        del tray_icon
+        try:
+            del tray_icon
+        except Exception as e:
+            print(f"Error destroying tray icon: {e}")
+    if mutex:
+        win32api.CloseHandle(mutex)
     exit(0)
 
-
 def load_icon(icon_path):
     try:
         return win32gui.LoadImage(None, icon_path, win32con.IMAGE_ICON, 0, 0, win32con.LR_LOADFROMFILE)
     except Exception as e:
         print(f"Error loading indicator icon: {e}")
         return None
 
-
 class SysTrayIcon:
     def __init__(self, hwnd, icon, tooltip):
         self.hwnd = hwnd
         self.icon = icon
         self.tooltip = tooltip
+        self.valid = True
         # System tray icon data structure
         self.nid = (
             self.hwnd,
@@ -110,9 +115,14 @@ class SysTrayIcon:
             self.tooltip,
         )
         # Add the icon to the system tray
-        win32gui.Shell_NotifyIcon(win32gui.NIM_ADD, self.nid)
+        try:
+            win32gui.Shell_NotifyIcon(win32gui.NIM_ADD, self.nid)
+        except Exception as e:
+            print(f"Error creating tray icon: {e}")
+            self.valid = False
 
     def set_tooltip(self, new_tooltip):
+        if not self.valid: return
         self.tooltip = new_tooltip
         self.nid = (
             self.hwnd,
@@ -122,9 +132,14 @@ class SysTrayIcon:
             self.icon,
             self.tooltip,
         )
-        win32gui.Shell_NotifyIcon(win32gui.NIM_MODIFY, self.nid)
+        try:
+            win32gui.Shell_NotifyIcon(win32gui.NIM_MODIFY, self.nid)
+        except Exception as e:
+            print(f"Error modifying tray icon tooltip: {e}")
+            self.valid = False
 
     def set_icon(self, icon):
+        if not self.valid: return
         self.icon = icon
         self.nid = (
             self.hwnd,
@@ -134,12 +149,20 @@ class SysTrayIcon:
             self.icon,
             self.tooltip,
         )
-        win32gui.Shell_NotifyIcon(win32gui.NIM_MODIFY, self.nid)
+        try:
+            win32gui.Shell_NotifyIcon(win32gui.NIM_MODIFY, self.nid)
+        except Exception as e:
+            print(f"Error modifying tray icon image: {e}")
+            self.valid = False
 
     def __del__(self):
         # Remove the icon from the system tray when the object is destroyed
-        win32gui.Shell_NotifyIcon(win32gui.NIM_DELETE, self.nid)
-
+        if not self.valid: return
+        try:
+            win32gui.Shell_NotifyIcon(win32gui.NIM_DELETE, self.nid)
+        except Exception as e:
+            print(f"Error deleting tray icon: {e}")
+            self.valid = False
 
 def get_current_color_prevalence():
     try:
@@ -158,7 +181,6 @@ def get_current_color_prevalence():
         print("Error accessing registry.")
         return None
 
-
 # Reads color palette binary data from a registry key and returns it as a hex string.
 def get_current_color_palette():
     try:
@@ -178,7 +200,6 @@ def get_current_color_palette():
         print("Error accessing registry.")
         return None
 
-
 # Attempts to get the current taskbar color based on user personalization settings and returns it in hex format.
 def get_current_taskbar_color():
     try:
@@ -203,7 +224,6 @@ def get_current_taskbar_color():
         print("Error accessing registry.")
         return None
 
-
 def set_taskbar_accent_color(hex_color, hex_color_palette, color_prevalence):
     """
     Sets the accent color location in the Windows registry.
@@ -248,7 +268,6 @@ def set_taskbar_accent_color(hex_color, hex_color_palette, color_prevalence):
     except WindowsError as e:
         print(f"Error accessing or modifying registry: {e}")
 
-
 def get_transparency_effects():
     try:
         key = winreg.OpenKey(
@@ -264,7 +283,6 @@ def get_transparency_effects():
     except WindowsError as e:
         print(f"Error accessing or modifying registry: {e}")
 
-
 def set_transparency_effects(value):
     try:
         key = winreg.OpenKey(
@@ -279,7 +297,6 @@ def set_transparency_effects(value):
     except WindowsError as e:
         print(f"Error accessing or modifying registry: {e}")
 
-
 # Attempt to extract a known good value in response.
 def extract_title(data):
     match = re.search(r"<title>(.*?)</title>", data)
@@ -288,7 +305,6 @@ def extract_title(data):
     else:
         return None
 
-
 def check_internet():
     for url, expected_response in CONNECT_TEST_URL_AND_RESPONSES.items():
         try:
@@ -301,7 +317,6 @@ def check_internet():
             pass
     return False
 
-
 def check_internet_and_update_tray_icon():
     global tray_icon, hicon_indicator_off, hicon_indicator_on, default_color
     if check_internet():
@@ -321,12 +336,28 @@ def check_internet_and_update_tray_icon():
         if get_wallpaper_path() != DEFAULT_BACKGROUND:  # Checked so program isn't continuously setting the wallpaper
             set_wallpaper(DEFAULT_BACKGROUND)
 
-
 def check_internet_loop():
+    global tray_icon
     while not stop_event.is_set():
-        check_internet_and_update_tray_icon()
-        time.sleep(CHECK_INTERVAL)
-
+        if tray_icon and tray_icon.valid:
+            check_internet_and_update_tray_icon()
+            time.sleep(CHECK_INTERVAL)
+        else:
+            print("Tray icon is invalid. Exiting check_internet_loop.")
+            stop_event.set()
+            if tray_icon:
+                del tray_icon
+                tray_icon = None
+
+            # Restart the tray icon and check_internet threads
+            tray_icon_thread = threading.Thread(target=tray_icon_loop)
+            tray_icon_thread.start()
+            # Wait for the tray icon to finish initializing
+            while tray_icon is None:
+                time.sleep(0.1)
+            check_thread = threading.Thread(target=check_internet_loop)
+            check_thread.start()
+            return
 
 def tray_icon_loop():
     global hwnd, tray_icon, hicon_indicator_off, hicon_indicator_on, stop_event
@@ -345,13 +376,13 @@ def tray_icon_loop():
     tray_icon = SysTrayIcon(hwnd, hicon_indicator_off, "Internet Detector")
 
     while not stop_event.is_set():
-        msg = win32gui.PeekMessage(hwnd, 0, 0, 0)
-        if msg and len(msg) == 6:
+        # Use PeekMessage to avoid blocking and allow thread exit
+        ret, msg = win32gui.PeekMessage(hwnd, 0, 0, win32con.PM_REMOVE)
+        if ret != 0:
             win32gui.TranslateMessage(msg)
             win32gui.DispatchMessage(msg)
         time.sleep(0.1)
 
-
 def get_wallpaper_path():
     """Attempts to retrieve the path to the current wallpaper image."""
     # Try to get the path from the registry (for wallpapers set through Windows settings)
@@ -378,7 +409,6 @@ def get_wallpaper_path():
     # If all else fails, return None
     return None
 
-
 def set_wallpaper(image_path):
     """Sets the desktop wallpaper to the image at the specified path."""
     print(f"Setting wallpaper to: {image_path}")
@@ -388,6 +418,32 @@ def set_wallpaper(image_path):
     if not result:
         print("Error setting wallpaper. Make sure the image path is correct.")
 
+def save_default_settings():
+    """Saves the default color, palette, and color prevalence to the registry."""
+    try:
+        key = winreg.CreateKey(winreg.HKEY_CURRENT_USER, REGISTRY_PATH)
+        winreg.SetValueEx(key, "DefaultColor", 0, winreg.REG_SZ, default_color)
+        winreg.SetValueEx(key, "DefaultPalette", 0, winreg.REG_SZ, default_palette)
+        winreg.SetValueEx(key, "DefaultColorPrevalence", 0, winreg.REG_DWORD, default_color_prevalence)
+        winreg.CloseKey(key)
+        print("Default settings saved to registry.")
+    except WindowsError as e:
+        print(f"Error saving default settings to registry: {e}")
+
+def load_default_settings():
+    """Loads the default color, palette, and color prevalence from the registry."""
+    global default_color, default_palette, default_color_prevalence
+    try:
+        key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, REGISTRY_PATH, 0, winreg.KEY_READ)
+        default_color, _ = winreg.QueryValueEx(key, "DefaultColor")
+        default_palette, _ = winreg.QueryValueEx(key, "DefaultPalette")
+        default_color_prevalence, _ = winreg.QueryValueEx(key, "DefaultColorPrevalence")
+        winreg.CloseKey(key)
+        print("Default settings loaded from registry.")
+        return True
+    except WindowsError:
+        print("No saved default settings found in registry.")
+        return False
 
 def main_loop():
     global stop_event, check_thread, tray_icon_thread, tray_icon, mutex
@@ -410,14 +466,18 @@ def main_loop():
     while not stop_event.is_set():
         time.sleep(1)
 
-
 if __name__ == "__main__":
     signal.signal(signal.SIGINT, signal_handler)
     urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
     default_transparency = get_transparency_effects()
-    default_color_prevalence = get_current_color_prevalence()
-    default_color = get_current_taskbar_color()
-    default_palette = get_current_color_palette()
+
+    # Try to load default settings from the registry
+    if not load_default_settings():
+        # If not found, get the current settings and save them as defaults
+        default_color_prevalence = get_current_color_prevalence()
+        default_color = get_current_taskbar_color()
+        default_palette = get_current_color_palette()
+        save_default_settings()
 
     # Create a hidden window to receive messages (required for system tray icons)
     def wndProc(hwnd, msg, wparam, lparam):
@@ -440,8 +500,4 @@ if __name__ == "__main__":
     print(f"Current wallpaper: {get_wallpaper_path()}")
     print(f"Current color: {default_color}")
 
-    main_loop()
-
-    # Release the mutex when the application exits
-    if mutex:
-        win32api.CloseHandle(mutex)
+    main_loop()