Improved aliases check

Author: zooba

_msbuild.py

@@ -84,6 +84,7 @@ class ResourceFile(CSourceFile):
     CFunction('date_as_str'),
     CFunction('datetime_as_str'),
     CFunction('reg_rename_key'),
+    CFunction('get_current_package'),
     CFunction('read_alias_package'),
     source='src/_native',
     RootNamespace='_native',

_msbuild_test.py

@@ -50,6 +50,7 @@
         CFunction('date_as_str'),
         CFunction('datetime_as_str'),
         CFunction('reg_rename_key'),
+        CFunction('get_current_package'),
         CFunction('read_alias_package'),
         source='src/_native',
     ),

src/_native/misc.cpp

@@ -1,5 +1,6 @@
 #include <Python.h>
 #include <windows.h>
+#include <appmodel.h>
 
 #include "helpers.h"
 
@@ -120,6 +121,22 @@ PyObject *reg_rename_key(PyObject *, PyObject *args, PyObject *kwargs) {
 }
 
 
+PyObject *get_current_package(PyObject *, PyObject *, PyObject *) {
+    wchar_t package_name[256];
+    UINT32 cch = sizeof(package_name) / sizeof(package_name[0]);
+    int err = GetCurrentPackageFamilyName(&cch, package_name);
+    switch (err) {
+    case ERROR_SUCCESS:
+        return PyUnicode_FromWideChar(package_name, cch);
+    case APPMODEL_ERROR_NO_PACKAGE:
+        return Py_GetConstant(Py_CONSTANT_NONE);
+    default:
+        PyErr_SetFromWindowsErr(err);
+        return NULL;
+    }
+}
+
+
 PyObject *read_alias_package(PyObject *, PyObject *args, PyObject *kwargs) {
     static const char * keywords[] = {"path", NULL};
     wchar_t *path = NULL;
@@ -136,22 +153,32 @@ PyObject *read_alias_package(PyObject *, PyObject *args, PyObject *kwargs) {
         return NULL;
     }
 
-    wchar_t buffer[32768];
+    struct {
+        DWORD tag;
+        DWORD _reserved1;
+        DWORD _reserved2;
+        wchar_t package_name[256];
+        wchar_t nul;
+    } buffer;
     DWORD nread;
 
     if (!DeviceIoControl(h, FSCTL_GET_REPARSE_POINT, NULL, 0,
-        buffer, sizeof(buffer), &nread, NULL)) {
+        &buffer, sizeof(buffer), &nread, NULL)
+        // we expect our buffer to be too small, but we only want the package
+        && GetLastError() != ERROR_MORE_DATA) {
         PyErr_SetFromWindowsErr(0);
         CloseHandle(h);
         return NULL;
     }
+
     CloseHandle(h);
 
-    if (*(DWORD*)buffer != IO_REPARSE_TAG_APPEXECLINK) {
+    if (buffer.tag != IO_REPARSE_TAG_APPEXECLINK) {
         return Py_GetConstant(Py_CONSTANT_NONE);
     }
 
-    return PyUnicode_FromWideChar(&buffer[4], nread / sizeof(wchar_t) - 5);
+    buffer.nul = 0;
+    return PyUnicode_FromWideChar(buffer.package_name, -1);
 }
 
 }

src/manage/firstrun.py

@@ -19,23 +19,36 @@
 
 def check_app_alias(cmd):
     LOGGER.debug("Checking app execution aliases")
-    from _native import read_alias_package
+    from _native import get_current_package, read_alias_package
+    # Expected identities:
+    # Side-loaded MSIX
+    # * "PythonSoftwareFoundation.PythonManager_3847v3x7pw1km",
+    # Store package
+    # * "PythonSoftwareFoundation.PythonManager_qbz5n2kfra8p0",
+    # Development build
+    # * "PythonSoftwareFoundation.PythonManager_m8z88z54g2w36",
+    # MSI/dev install
+    # * None
+    try:
+        pkg = get_current_package()
+    except OSError:
+        LOGGER.debug("Failed to get current package name.", exc_info=True)
+        pkg = None
+    if not pkg:
+        pkg = ""
+        #LOGGER.debug("Check skipped: MSI install can't do this check")
+        #return True
+    LOGGER.debug("Checking for %s", pkg)
     root = Path(os.environ["LocalAppData"]) / "Microsoft/WindowsApps"
     for name in ["py.exe", "pyw.exe", "python.exe", "pythonw.exe", "python3.exe", "pymanager.exe"]:
         exe = root / name
         try:
             LOGGER.debug("Reading from %s", exe)
-            package = (read_alias_package(exe) or "").split("\0")
-            LOGGER.debug("Data: %r", package)
-            if package[1] not in (
-                # Side-loaded MSIX
-                "PythonSoftwareFoundation.PythonManager_3847v3x7pw1km",
-                # Store packaged
-                "PythonSoftwareFoundation.PythonManager_qbz5n2kfra8p0",
-                # Development build
-                "PythonSoftwareFoundation.PythonManager_m8z88z54g2w36",
-            ):
+            package = read_alias_package(exe)
+            LOGGER.debug("Package: %r", package)
+            if package not in pkg:
                 LOGGER.debug("Check failed: package did not match identity")
+                return False
         except FileNotFoundError:
             LOGGER.debug("Check failed: did not find %s", exe)
             return False
@@ -59,7 +72,10 @@ def check_long_paths(cmd):
 
 def check_py_on_path(cmd):
     LOGGER.debug("Checking for legacy py.exe on PATH")
-    from _native import read_alias_package
+    from _native import get_current_package, read_alias_package
+    if not get_current_package():
+        LOGGER.debug("Check skipped: MSI install can't do this check")
+        return True
     for p in os.environ["PATH"].split(";"):
         if not p:
             continue
@@ -230,7 +246,7 @@ def first_run(cmd):
 if __name__ == "__main__":
     class TestCommand:
         enabled = True
-        global_dir = r".\test-bin"
+        global_dir = Path(os.path.expandvars(r"%LocalAppData%\Python\bin"))
         explicit = False
         confirm = True
         check_app_alias = True
@@ -241,7 +257,16 @@ class TestCommand:
         check_default_tag = True
 
         def get_installs(self, *args, **kwargs):
-            return []
+            import json
+            root = Path(os.path.expandvars(r"%LocalAppData%\Python"))
+            result = []
+            for d in root.iterdir():
+                inst = d / "__install__.json"
+                try:
+                    result.append(json.loads(inst.read_text()))
+                except FileNotFoundError:
+                    pass
+            return result
 
         def _ask(self, fmt, *args, yn_text="Y/n", expect_char="y"):
             if not self.confirm: