Add internet detector utility

scripts/utils/internet_detect_tool.py

@@ -0,0 +1,246 @@
+from PyQt5.QtWidgets import QApplication, QMenu
+from PyQt5.QtCore import QTimer, QThread, pyqtSignal
+import threading
+import requests
+import urllib3
+import win32api
+import win32gui
+import win32con
+import winreg
+import struct
+import signal
+import ctypes
+import time
+import os
+
+
+# Define constants
+CHECK_INTERVAL = 2  # Seconds
+CONNECT_TEST_URL = "https://www.msftconnecttest.com/connecttest.txt"
+EXPECTED_RESPONSE = "Microsoft Connect Test"
+SPI_SETDESKWALLPAPER = 20
+SPIF_UPDATEINIFILE = 0x01
+SPIF_SENDWININICHANGE = 0x02
+COLOR_DESKTOP = 1
+ICON_RED_X = "red_circle.ico"
+ICON_FLARE = "flare.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")
+
+# Global variables
+app = QApplication([])
+tray_icon = None
+stop_event = threading.Event()  # To signal the background thread to stop
+hwnd = None     # We'll assign the window handle here later
+# Win32 API icon handles
+hicon_flare = None
+hicon_red_x = None
+
+def signal_handler(sig, frame):
+    print("Ctrl+C detected. Resetting background and exiting...")
+    reset_background()
+    stop_event.set()    # Signal the background thread to stop
+    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 flare icon: {e}")
+        return None
+
+class SysTrayIcon:
+    def __init__(self, hwnd, icon, tooltip):
+        self.hwnd = hwnd
+        self.icon = icon
+        self.tooltip = tooltip
+
+        # System tray icon data structure
+        self.nid = (self.hwnd, 0, win32gui.NIF_ICON | win32gui.NIF_MESSAGE | win32gui.NIF_TIP, 
+                    win32con.WM_USER + 20, self.icon, self.tooltip)
+
+        # Add the icon to the system tray
+        win32gui.Shell_NotifyIcon(win32gui.NIM_ADD, self.nid)
+
+        # Create a context menu
+        self.menu = win32gui.CreatePopupMenu()
+        win32gui.AppendMenu(self.menu, win32con.MF_STRING, 1023, "Exit")  # Example menu item
+
+    def show_context_menu(self, x, y):
+        # Display the context menu at the specified coordinates
+        win32gui.SetForegroundWindow(self.hwnd)
+        win32gui.TrackPopupMenu(self.menu, win32con.TPM_LEFTALIGN, x, y, 0, self.hwnd, None)
+        win32gui.PostMessage(self.hwnd, win32con.WM_NULL, 0, 0)
+
+    def set_icon(self, icon):
+        self.icon = icon
+        self.nid = (self.hwnd, 0, win32gui.NIF_ICON | win32gui.NIF_MESSAGE | win32gui.NIF_TIP, 
+                    win32con.WM_USER + 20, self.icon, self.tooltip)
+        win32gui.Shell_NotifyIcon(win32gui.NIM_MODIFY, self.nid)
+
+    def show_balloon_tip(self, title, msg):
+        # Display a balloon tip notification
+        win32gui.Shell_NotifyIcon(win32gui.NIM_MODIFY, (
+            self.hwnd,
+            0,
+            win32gui.NIF_INFO,
+            win32con.WM_USER + 20,
+            self.icon,
+            self.tooltip,
+            msg,
+            200,
+            title,
+        ))
+
+    def __del__(self):
+        # Remove the icon from the system tray when the object is destroyed
+        win32gui.Shell_NotifyIcon(win32gui.NIM_DELETE, self.nid)
+
+class CheckInternetThread(QThread):
+    update_signal = pyqtSignal()
+
+    def run(self):
+        while not stop_event.is_set():
+            self.update_signal.emit()
+            time.sleep(CHECK_INTERVAL)
+
+class TrayIconThread(QThread):
+    icon_created = pyqtSignal() # Signal to indicate icon creation
+    def __init__(self, parent=None):
+        super().__init__(parent)
+        self.tray_icon = None
+
+    def run(self):
+        global hwnd, hicon_flare, hicon_red_x
+        # Load icons
+        hicon_flare = load_icon(ICON_FLARE)
+        hicon_red_x = load_icon(ICON_RED_X)
+
+        # Wait for hwnd to be initialized
+        while hwnd is None:
+            time.sleep(0.1)
+
+        if hicon_flare is None or hicon_red_x is None:
+            print("Error: Failed to load icons. Exiting TrayIconThread.")
+            return
+
+        # Create the system tray icon
+        self.tray_icon = SysTrayIcon(hwnd, hicon_flare, "FLARE Internet Detector")
+
+        win32gui.PumpMessages() # Continuously checks the Windows message queue for messages sent to application's window
+
+def enable_transparency_effects():
+  try:
+      key = winreg.OpenKey(winreg.HKEY_CURRENT_USER,
+                           r"Software\Microsoft\Windows\CurrentVersion\Themes\Personalize",
+                           0, winreg.KEY_ALL_ACCESS)
+
+      winreg.SetValueEx(key, "EnableTransparency", 0, winreg.REG_DWORD, 1)
+      winreg.CloseKey(key)
+  except WindowsError as e:
+      print(f"Error accessing or modifying registry: {e}")
+
+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)
+    try:
+        key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, r"Control Panel\Desktop", 0, winreg.KEY_READ)
+        value, _ = winreg.QueryValueEx(key, "Wallpaper")
+        winreg.CloseKey(key)
+        if value:
+            return value
+    except WindowsError:
+        pass
+
+    # Check for cached wallpaper files (if the above fails)
+    cached_files_dir = os.path.join(os.getenv("APPDATA"), r"Microsoft\Windows\Themes\CachedFiles")
+    transcoded_wallpaper_path = os.path.join(os.getenv("APPDATA"), r"Microsoft\Windows\Themes\TranscodedWallpaper")
+
+    for file in os.listdir(cached_files_dir):
+        if file.endswith(('.jpg', '.jpeg', '.bmp', '.png')):
+            return os.path.join(cached_files_dir, file)
+
+    if os.path.exists(transcoded_wallpaper_path):
+        return transcoded_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("Setting wallpaper to: {}".format(image_path))
+    result = ctypes.windll.user32.SystemParametersInfoW(
+        SPI_SETDESKWALLPAPER, 0, image_path, SPIF_UPDATEINIFILE | SPIF_SENDWININICHANGE
+    )
+    if not result:
+        print("Error setting wallpaper. Make sure the image path is correct.")
+
+def reset_background():
+    # Set background back to default wallpaper
+    set_wallpaper(DEFAULT_BACKGROUND)
+
+def check_internet():
+    try:
+        response = requests.get(CONNECT_TEST_URL, timeout=5, verify=False)
+        return EXPECTED_RESPONSE in response.text
+    except requests.exceptions.RequestException:
+        return False
+
+def update_tray_icon():
+    global tray_icon, hicon_flare, hicon_red_x
+    if check_internet():
+        tray_icon.set_icon(hicon_red_x)
+        # Set the background to internet connection background
+        if get_wallpaper_path() != INTERNET_BACKGROUND:     # Checked so program isn't continuously setting the wallpaper
+            set_wallpaper(INTERNET_BACKGROUND)
+    else:  # Reset background when internet is not detected
+        tray_icon.set_icon(hicon_flare)
+        # Reset background to default
+        if get_wallpaper_path() != DEFAULT_BACKGROUND:      # Checked so program isn't continuously setting the wallpaper
+            reset_background()
+
+def main_loop():
+    # Create and start the threads
+    check_thread = CheckInternetThread()
+    tray_icon_thread = TrayIconThread()
+    check_thread.update_signal.connect(update_tray_icon)
+    check_thread.start()
+    tray_icon_thread.start()
+
+    # Wait for the tray icon to finish initializing
+    while tray_icon_thread.tray_icon is None:
+        time.sleep(0.1)
+
+    global tray_icon
+    tray_icon = tray_icon_thread.tray_icon  # Get the tray icon object from the thread
+
+    app.exec_()
+
+
+
+if __name__ == "__main__":
+    signal.signal(signal.SIGINT, signal_handler)
+    urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+    enable_transparency_effects()
+
+    # Create a hidden window to receive messages (required for system tray icons)
+    def wndProc(hwnd, msg, wparam, lparam):
+        if lparam == win32con.WM_LBUTTONDBLCLK:
+            print("Left button double clicked")
+        elif msg == win32con.WM_COMMAND:
+            if wparam == 1023:  # Example menu item ID
+                print("Exit selected")
+                win32gui.DestroyWindow(hwnd)
+                app.quit()  # Exit the PyQt application
+        return win32gui.DefWindowProc(hwnd, msg, wparam, lparam)
+
+    wc = win32gui.WNDCLASS()
+    hinst = wc.hInstance = win32api.GetModuleHandle(None)
+    wc.lpszClassName = "FLARE Internet Detector"
+    wc.lpfnWndProc = wndProc
+    classAtom = win32gui.RegisterClass(wc)
+    hwnd = win32gui.CreateWindow(classAtom, "FLARE Internet Detector", 0, 0, 0, 0, 0, 0, 0, hinst, None)
+
+    print("Current wallpaper: {}".format(DEFAULT_BACKGROUND))
+
+    main_loop()