From d42784446518a9e3796573d9564f9d2f9728ceae Mon Sep 17 00:00:00 2001
From: peace-maker <peacemakerctf@gmail.com>
Date: Thu, 18 Jan 2024 00:57:43 +0100
Subject: [PATCH] Lookup using $PATHEXT file extensions in `which` on Windows
 (#2328)

* Lookup process executable using PATHEXT file extensions

Windows uses some file extensions to determine which file to run when given
a non-existing file.

This allows to run `calc` instead of `calc.exe`.

* Try PATHEXT in `which` directly

* Update CHANGELOG
---
 CHANGELOG.md            |  2 ++
 pwnlib/tubes/process.py |  8 ++++++--
 pwnlib/util/misc.py     | 40 +++++++++++++++++++++++++---------------
 3 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 51cfab80a..b0391041f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -81,6 +81,7 @@ The table below shows which release corresponds to each branch, and what date th
 - [#2279][2279] Make `pwn template` always set context.binary
 - [#2310][2310] Add support to start a process on Windows
 - [#2334][2334] Speed up disasm commandline tool with colored output
+- [#2328][2328] Lookup using $PATHEXT file extensions in `which` on Windows
 
 [2242]: https://github.com/Gallopsled/pwntools/pull/2242
 [2277]: https://github.com/Gallopsled/pwntools/pull/2277
@@ -93,6 +94,7 @@ The table below shows which release corresponds to each branch, and what date th
 [2279]: https://github.com/Gallopsled/pwntools/pull/2279
 [2310]: https://github.com/Gallopsled/pwntools/pull/2310
 [2334]: https://github.com/Gallopsled/pwntools/pull/2334
+[2328]: https://github.com/Gallopsled/pwntools/pull/2328
 
 ## 4.12.0 (`beta`)
 
diff --git a/pwnlib/tubes/process.py b/pwnlib/tubes/process.py
index af9879bd9..e5f143e75 100644
--- a/pwnlib/tubes/process.py
+++ b/pwnlib/tubes/process.py
@@ -33,7 +33,7 @@
 from pwnlib.util.misc import parse_ldd_output
 from pwnlib.util.misc import which
 from pwnlib.util.misc import normalize_argv_env
-from pwnlib.util.packing import _need_bytes
+from pwnlib.util.packing import _decode
 
 log = getLogger(__name__)
 
@@ -558,7 +558,11 @@ def _validate(self, cwd, executable, argv, env):
 
         argv, env = normalize_argv_env(argv, env, self, 4)
         if env:
-            env = {bytes(k): bytes(v) for k, v in env}
+            if sys.platform == 'win32':
+                # Windows requires that all environment variables be strings
+                env = {_decode(k): _decode(v) for k, v in env}
+            else:
+                env = {bytes(k): bytes(v) for k, v in env}
         if argv:
             argv = list(map(bytes, argv))
 
diff --git a/pwnlib/util/misc.py b/pwnlib/util/misc.py
index 2ca9f31a5..691539471 100644
--- a/pwnlib/util/misc.py
+++ b/pwnlib/util/misc.py
@@ -139,6 +139,7 @@ def which(name, all = False, path=None):
 
     Works as the system command ``which``; searches $PATH for ``name`` and
     returns a full path if found.
+    Tries all of the file extensions in $PATHEXT on Windows too.
 
     If `all` is :const:`True` the set of all found locations is returned, else
     the first occurrence or :const:`None` is returned.
@@ -160,26 +161,35 @@ def which(name, all = False, path=None):
     if os.path.sep in name:
         return name
 
-    isroot = False if sys.platform == 'win32' else (os.getuid() == 0)
+    if sys.platform == 'win32':
+        pathexts = os.environ.get('PATHEXT', '').split(os.pathsep)
+        isroot = False
+    else:
+        pathexts = []
+        isroot = os.getuid() == 0
+    pathexts = [''] + pathexts
     out = set()
     try:
         path = path or os.environ['PATH']
     except KeyError:
         log.exception('Environment variable $PATH is not set')
-    for p in path.split(os.pathsep):
-        p = os.path.join(p, name)
-        if os.access(p, os.X_OK):
-            st = os.stat(p)
-            if not stat.S_ISREG(st.st_mode):
-                continue
-            # work around this issue: https://bugs.python.org/issue9311
-            if isroot and not \
-              st.st_mode & (stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH):
-                continue
-            if all:
-                out.add(p)
-            else:
-                return p
+    for path_part in path.split(os.pathsep):
+        for ext in pathexts:
+            nameext = name + ext
+            p = os.path.join(path_part, nameext)
+            if os.access(p, os.X_OK):
+                st = os.stat(p)
+                if not stat.S_ISREG(st.st_mode):
+                    continue
+                # work around this issue: https://bugs.python.org/issue9311
+                if isroot and not \
+                st.st_mode & (stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH):
+                    continue
+                if all:
+                    out.add(p)
+                    break
+                else:
+                    return p
     if all:
         return out
     else: