flares/file_ext_registration.py

"""
Filename: file_ext_registration.py
Description: Registers the extension of project files with the application
Note: Compliant with pylance strict type checking

Author: Tyler de Zeeuw
License: GPL-3.0
"""

# Built-in imports
import os
import sys
import plistlib
import subprocess
from typing import Optional, Tuple

# External library imports
from src.shared.shareddata import APP_NAME, PLATFORM_NAME

ELEVATION_FLAG = "--register_file_association_elevated"



def register_file_association(ext: Optional[str] = None,
    prog_id: Optional[str] = None,
    app_name: Optional[str] = None,
    bundle_id: Optional[str] = None,
    force_admin: bool = False,
) -> Tuple[bool, str]:
    
    """
    Registers a custom file extension across Windows, Linux, and macOS.
    Handles non-admin Windows users by falling back to local user registry.
    """

    clean_app = APP_NAME.replace(" ", "").lower()[:-1]
    if ext is None:
        ext = f".{clean_app}"
    if prog_id is None:
        prog_id = f"{APP_NAME.upper().replace(' ', '')}.ProjectFile"
    if app_name is None:
        app_name = APP_NAME

    # Ensure extension starts with dot and contains no spaces
    ext = f".{ext.lstrip('.').replace(' ', '').lower()}"

    if PLATFORM_NAME == "windows":
        return _register_windows(ext, prog_id, app_name, force_admin=force_admin)
    elif PLATFORM_NAME == "linux":
        return _register_linux(ext, prog_id, app_name)
    elif PLATFORM_NAME == "darwin":  # macOS
        if bundle_id is None:
            bundle_id = f"com.{clean_app}.app"
        return _register_macos(ext, app_name, bundle_id)
    else:
        return False, f"Unsupported OS: {PLATFORM_NAME}"



def _dev_windowless_executable() -> str:
    """
    Windows dev-mode only: returns pythonw.exe alongside the current
    interpreter if it exists, otherwise falls back to sys.executable.
    """

    if getattr(sys, 'frozen', False):
        return sys.executable
    exe_dir = os.path.dirname(sys.executable)
    windowless = os.path.join(exe_dir, "pythonw.exe")
    return windowless if os.path.exists(windowless) else sys.executable



def is_windows_admin() -> bool:
    """
    Returns True if currently running elevated on Windows. Always False on
    macOS/Linux, and False (rather than raising) if the check itself fails.
    """

    if PLATFORM_NAME != "windows":
        return False
    import ctypes
    try:
        return ctypes.windll.shell32.IsUserAnAdmin() != 0
    except Exception:
        return False



def _relaunch_elevated(ext: str, prog_id: str, app_name: str) -> Optional[int]:
    """
    Triggers a UAC prompt and re-launches this process elevated with the
    registration args.
    """

    import ctypes
    from ctypes import wintypes

    if getattr(sys, 'frozen', False):
        exe = sys.executable
        params = [ELEVATION_FLAG, ext, prog_id, app_name]
    else:
        exe = _dev_windowless_executable()
        script_path = os.path.abspath(sys.argv[0])
        params = [script_path, ELEVATION_FLAG, ext, prog_id, app_name]

    param_str = " ".join(f'"{p}"' for p in params)

    SEE_MASK_NOCLOSEPROCESS = 0x00000040
    SW_SHOWNORMAL = 1

    class SHELLEXECUTEINFO(ctypes.Structure):
        _fields_ = [
            ("cbSize", wintypes.DWORD),
            ("fMask", ctypes.c_ulong),
            ("hwnd", wintypes.HWND),
            ("lpVerb", wintypes.LPCWSTR),
            ("lpFile", wintypes.LPCWSTR),
            ("lpParameters", wintypes.LPCWSTR),
            ("lpDirectory", wintypes.LPCWSTR),
            ("nShow", ctypes.c_int),
            ("hInstApp", wintypes.HINSTANCE),
            ("lpIDList", ctypes.c_void_p),
            ("lpClass", wintypes.LPCWSTR),
            ("hKeyClass", wintypes.HKEY),
            ("dwHotKey", wintypes.DWORD),
            ("hIconOrMonitor", wintypes.HANDLE),
            ("hProcess", wintypes.HANDLE),
        ]

    sei = SHELLEXECUTEINFO()
    sei.cbSize = ctypes.sizeof(sei)
    sei.fMask = SEE_MASK_NOCLOSEPROCESS
    sei.lpVerb = "runas"
    sei.lpFile = exe
    sei.lpParameters = param_str
    sei.nShow = SW_SHOWNORMAL

    if not ctypes.windll.shell32.ShellExecuteExW(ctypes.byref(sei)):
        # User clicked "No" on the UAC prompt, or elevation failed outright.
        return None

    WAIT_INFINITE = 0xFFFFFFFF
    ctypes.windll.kernel32.WaitForSingleObject(sei.hProcess, WAIT_INFINITE)

    exit_code = wintypes.DWORD()
    ctypes.windll.kernel32.GetExitCodeProcess(sei.hProcess, ctypes.byref(exit_code))
    ctypes.windll.kernel32.CloseHandle(sei.hProcess)

    return exit_code.value



def _delete_key_recursive(root_key: int, path: str) -> None:
    """Recursively deletes a registry key and all its subkeys, if present."""

    import winreg
    try:
        with winreg.OpenKey(root_key, path, 0, winreg.KEY_ALL_ACCESS) as key:
            while True:
                try:
                    subkey_name = winreg.EnumKey(key, 0)
                except OSError:
                    break
                _delete_key_recursive(root_key, f"{path}\\{subkey_name}")
        winreg.DeleteKey(root_key, path)
    except FileNotFoundError:
        pass



def _cleanup_stale_hkcu_entries(ext: str, prog_id: str) -> None:
    """
    Removes any leftover per-user (HKCU) association for this ext/prog_id.
    Windows prefers HKCU\\Software\\Classes over HKEY_CLASSES_ROOT for a
    given user, so a stale non-admin dev-time registration can silently
    keep shadowing a correct system-wide (admin) registration.
    """

    import winreg
    _delete_key_recursive(winreg.HKEY_CURRENT_USER, f"Software\\Classes\\{ext}")
    _delete_key_recursive(winreg.HKEY_CURRENT_USER, f"Software\\Classes\\{prog_id}")



def _register_windows(ext: str, prog_id: str, app_name: str, force_admin: bool = False) -> Tuple[bool, str]:
    """
    Registers the file extension with the application on Windows.
    """
    import winreg
    import ctypes

    is_admin = is_windows_admin()

    if force_admin and not is_admin:
        exit_code = _relaunch_elevated(ext, prog_id, app_name)
        if exit_code is None:
            return False, "Elevation was cancelled or the UAC prompt could not be shown."
        if exit_code != 0:
            return False, f"Elevated registration process exited with code {exit_code}."
        return True, f"Successfully registered {ext} on Windows (System-wide, via elevated relaunch)!"

    # Windows is the one platform where unfrozen invocation is supported
    if getattr(sys, 'frozen', False):
        command_str = f'"{sys.executable}" "%1"'
    else:
        script_path = os.path.abspath(sys.argv[0])
        command_str = f'"{_dev_windowless_executable()}" "{script_path}" "%1"'

    root_key = winreg.HKEY_CLASSES_ROOT if is_admin else winreg.HKEY_CURRENT_USER
    base_path = "" if is_admin else "Software\\Classes\\"

    try:
        # 1. Map extension -> ProgID
        ext_path = f"{base_path}{ext}"
        with winreg.CreateKey(root_key, ext_path) as key:
            winreg.SetValue(key, "", winreg.REG_SZ, prog_id)

        # 2. Add to OpenWithProgids (Forces Windows 10/11 to show in 'Open With')
        open_with_path = f"{base_path}{ext}\\OpenWithProgids"
        with winreg.CreateKey(root_key, open_with_path) as key:
            winreg.SetValueEx(key, prog_id, 0, winreg.REG_SZ, "")

        # 3. Set friendly type description (shown as "Type" in File Explorer)
        type_description = f"{app_name.upper()} Project File"
        prog_path = f"{base_path}{prog_id}"
        with winreg.CreateKey(root_key, prog_path) as key:
            winreg.SetValue(key, "", winreg.REG_SZ, type_description)

        # 4. Set launch command
        cmd_path = f"{base_path}{prog_id}\\shell\\open\\command"
        with winreg.CreateKey(root_key, cmd_path) as key:
            winreg.SetValue(key, "", winreg.REG_SZ, command_str)

        # 5. Notify Shell (SHCNE_ASSOCCHANGED = 0x08000000, SHCNF_IDLIST = 0x0000)
        try:
            ctypes.windll.shell32.SHChangeNotify(0x08000000, 0x0000, 0, 0)
        except Exception:
            pass

        # 6. If we just wrote a system-wide entry, clean up any stale per-user entry that would otherwise shadow it.
        if is_admin:
            try:
                _cleanup_stale_hkcu_entries(ext, prog_id)
            except Exception:
                # Non-fatal: the system-wide write above already succeeded.
                pass

        scope = "System-wide" if is_admin else "Local User (Non-Admin)"
        return True, f"Successfully registered {ext} on Windows ({scope})!"

    except Exception as e:
        return False, f"Windows Registration failed: {str(e)}"


# TODO: Validate
def _register_linux(ext: str, prog_id: str, app_name: str) -> Tuple[bool, str]:
    """
    Registers the file extension with the application on Linux based platforms.
    """
    clean_ext = ext.lstrip('.')
    mime_type = f"application/x-{clean_ext}"
    desktop_file_name = f"{prog_id.lower()}.desktop"

    # Linux build is always a frozen PyInstaller executable
    exe_path = f'"{sys.executable}"'

    apps_dir = os.path.expanduser("~/.local/share/applications")
    mime_dir = os.path.expanduser("~/.local/share/mime/packages")
    os.makedirs(apps_dir, exist_ok=True)
    os.makedirs(mime_dir, exist_ok=True)

    try:
        # 1. Create XML MIME definition
        mime_xml_path = os.path.join(mime_dir, f"{prog_id.lower()}.xml")
        xml_content = f"""<?xml version="1.0" encoding="UTF-8"?>
            <mime-info xmlns="http://www.freedesktop.org/standards/shared-mime-info">
            <mime-type type="{mime_type}">
                <comment>{app_name} File</comment>
                <glob pattern="*.{clean_ext}"/>
            </mime-type>
            </mime-info>"""
        with open(mime_xml_path, "w") as f:
            f.write(xml_content)

        # 2. Create .desktop file
        desktop_path = os.path.join(apps_dir, desktop_file_name)
        desktop_content = f"""[Desktop Entry]
            Name={app_name}
            Exec={exe_path} %f
            Type=Application
            MimeType={mime_type};
            Terminal=false
            """
        with open(desktop_path, "w") as f:
            f.write(desktop_content)

        # 3. Update the shared-mime-info database and check it actually worked
        mime_result = subprocess.run(
            ["update-mime-database", os.path.expanduser("~/.local/share/mime")],
            capture_output=True, text=True,
        )
        if mime_result.returncode != 0:
            return False, f"update-mime-database failed: {mime_result.stderr.strip()}"

        # 4. Set as default handler and check it actually worked
        xdg_result = subprocess.run(
            ["xdg-mime", "default", desktop_file_name, mime_type],
            capture_output=True, text=True,
        )
        if xdg_result.returncode != 0:
            return False, f"xdg-mime default failed: {xdg_result.stderr.strip()}"

        return True, f"Successfully registered {ext} on Linux (Local User)!"

    except FileNotFoundError as e:
        return False, (
            f"Required tool not found ({str(e)}). Is shared-mime-info / "
            f"xdg-utils installed on this system?"
        )
    except Exception as e:
        return False, f"Linux Registration failed: {str(e)}"


# TODO: Validate
def _register_macos(ext: str, app_name: str, bundle_id: str) -> Tuple[bool, str]:
    """
    Registers the file extension with the application on macOS.
    """

    if not getattr(sys, 'frozen', False):
        return False, "macOS file association requires app to be packaged as a .app bundle."

    exe_path = sys.executable
    app_bundle_path = os.path.abspath(os.path.join(exe_path, "../../../"))

    if not app_bundle_path.endswith(".app"):
        return False, "Could not locate outer .app bundle."

    info_plist_path = os.path.join(app_bundle_path, "Contents", "Info.plist")
    if not os.path.exists(info_plist_path):
        return False, f"Info.plist not found at {info_plist_path}"

    clean_ext = ext.lstrip('.')
    uti = f"{bundle_id}.{clean_ext}"

    try:
        with open(info_plist_path, "rb") as f:
            plist = plistlib.load(f)

        # Tell macOS this app can open files with our UTI
        doc_types = plist.get("CFBundleDocumentTypes", [])
        if not any(uti in dt.get("LSItemContentTypes", []) for dt in doc_types):
            doc_types.append({
                "CFBundleTypeName": f"{app_name} Project File",
                "CFBundleTypeRole": "Editor",
                "LSHandlerRank": "Owner",
                "LSItemContentTypes": [uti],
            })
        plist["CFBundleDocumentTypes"] = doc_types

        # Declare the UTI itself (required — without this the type is unknown to LS)
        exported_types = plist.get("UTExportedTypeDeclarations", [])
        if not any(t.get("UTTypeIdentifier") == uti for t in exported_types):
            exported_types.append({
                "UTTypeIdentifier": uti,
                "UTTypeDescription": f"{app_name} File",
                "UTTypeConformsTo": ["public.data"],
                "UTTypeTagSpecification": {"public.filename-extension": [clean_ext]},
            })
        plist["UTExportedTypeDeclarations"] = exported_types

        with open(info_plist_path, "wb") as f:
            plistlib.dump(plist, f)

    except Exception as e:
        return False, f"Failed to update Info.plist: {str(e)}"

    try:
        lsregister_path = (
            "/System/Library/Frameworks/CoreServices.framework/Frameworks/"
            "LaunchServices.framework/Support/lsregister"
        )
        subprocess.run(
            [lsregister_path, "-f", app_bundle_path],
            check=True, capture_output=True,
        )
        return True, f"Registered {ext} with {app_bundle_path} via macOS Launch Services!"
    except subprocess.CalledProcessError as e:
        stderr = e.stderr.decode(errors="ignore") if e.stderr else str(e)
        return False, f"macOS Launch Services registration failed: {stderr}"
    except Exception as e:
        return False, f"macOS Registration failed: {str(e)}"