fix: traceback file maybe use local import for trace print

Signed-off-by: yihong0618 <[email protected]>

Author: yihong0618

Include/internal/pycore_traceback.h

@@ -14,6 +14,10 @@ PyAPI_FUNC(int) _Py_DisplaySourceLine(PyObject *, PyObject *, int, int, int *, P
 // Export for 'pyexact' shared extension
 PyAPI_FUNC(void) _PyTraceback_Add(const char *, const char *, int);
 
+// Helper function to check if it's safe to import traceback module
+// Returns 0 if a traceback.py file exists in sys.path[0], 1 if safe
+extern int _PyTraceback_IsSafeToImport(void);
+
 /* Write the Python traceback into the file 'fd'. For example:
 
        Traceback (most recent call first):

Lib/test/test_traceback.py

@@ -5175,6 +5175,43 @@ def expected(t, m, fn, l, f, E, e, z):
                 f'    +------------------------------------',
         ]
         self.assertEqual(actual, expected(**colors))
+class TestTracebackSafety(unittest.TestCase):
+    """Test security fixes for traceback module shadowing (issue gh-138170)"""
+
+    def setUp(self):
+        self.tempdir = tempfile.mkdtemp()
+        self.addCleanup(shutil.rmtree, self.tempdir)
+
+    def test_traceback_shadowing_protection(self):
+        """Test that traceback.py files are not executed during exception handling"""
+        import os
+        import sys
+        import subprocess
+        malicious_traceback = os.path.join(self.tempdir, 'traceback.py')
+        with open(malicious_traceback, 'w') as f:
+            f.write('''
+print("MALICIOUS CODE EXECUTED - SECURITY BREACH!")
+print("This should not appear in traceback output")
+''')
+        test_script = os.path.join(self.tempdir, 'test_exception.py')
+        with open(test_script, 'w') as f:
+            f.write('''
+import sys
+sys.path.insert(0, ".")
+raise ValueError("test exception")
+''')
+        old_cwd = os.getcwd()
+        try:
+            os.chdir(self.tempdir)
+            result = subprocess.run([sys.executable, test_script],
+                                  capture_output=True, text=True)
+            self.assertIn("ValueError: test exception", result.stderr)
+            self.assertNotIn("MALICIOUS CODE EXECUTED", result.stdout)
+            self.assertNotIn("MALICIOUS CODE EXECUTED", result.stderr)
+            self.assertNotIn("This should not appear in traceback output", result.stdout)
+            self.assertNotIn("This should not appear in traceback output", result.stderr)
+        finally:
+            os.chdir(old_cwd)
 
 if __name__ == "__main__":
     unittest.main()

Python/errors.c

@@ -13,6 +13,10 @@
 #include "pycore_traceback.h"     // _PyTraceBack_FromFrame()
 #include "pycore_unicodeobject.h" // _PyUnicode_Equal()
 
+#include <stdio.h>                // fopen, fclose
+#include <limits.h>               // PATH_MAX
+#include <string.h>               // strlen
+
 #ifdef MS_WINDOWS
 #  include <windows.h>
 #  include <winbase.h>
@@ -1488,25 +1492,27 @@ write_unraisable_exc_file(PyThreadState *tstate, PyObject *exc_type,
         }
     }
 
-    // Try printing the exception using the stdlib module.
-    // If this fails, then we have to use the C implementation.
-    PyObject *print_exception_fn = PyImport_ImportModuleAttrString("traceback",
-                                                                   "_print_exception_bltin");
-    if (print_exception_fn != NULL && PyCallable_Check(print_exception_fn)) {
-        PyObject *args[2] = {exc_value, file};
-        PyObject *result = PyObject_Vectorcall(print_exception_fn, args, 2, NULL);
-        int ok = (result != NULL);
-        Py_DECREF(print_exception_fn);
-        Py_XDECREF(result);
-        if (ok) {
-            // Nothing else to do
-            return 0;
+    // Try printing the exception using the stdlib module, but be careful about
+    // traceback module shadowing (issue gh-138170). Use safe import check.
+    if (_PyTraceback_IsSafeToImport()) {
+        PyObject *print_exception_fn = PyImport_ImportModuleAttrString("traceback", "_print_exception_bltin");
+        if (print_exception_fn != NULL && PyCallable_Check(print_exception_fn)) {
+            PyObject *args[2] = {exc_value, file};
+            PyObject *result = PyObject_Vectorcall(print_exception_fn, args, 2, NULL);
+            int ok = (result != NULL);
+            Py_DECREF(print_exception_fn);
+            Py_XDECREF(result);
+            if (ok) {
+                // Nothing else to do
+                return 0;
+            }
+        }
+        else {
+            Py_XDECREF(print_exception_fn);
         }
     }
-    else {
-        Py_XDECREF(print_exception_fn);
-    }
-    // traceback module failed, fall back to pure C
+
+    // traceback module failed or unsafe, fall back to pure C
     _PyErr_Clear(tstate);
 
     if (exc_tb != NULL && exc_tb != Py_None) {

Python/pythonrun.c

@@ -1146,20 +1146,21 @@ _PyErr_Display(PyObject *file, PyObject *unused, PyObject *value, PyObject *tb)
         }
     }
 
-    // Try first with the stdlib traceback module
-    PyObject *print_exception_fn = PyImport_ImportModuleAttrString(
-        "traceback",
-        "_print_exception_bltin");
-    if (print_exception_fn == NULL || !PyCallable_Check(print_exception_fn)) {
-        goto fallback;
-    }
-
-    PyObject* result = PyObject_CallOneArg(print_exception_fn, value);
-
-    Py_XDECREF(print_exception_fn);
-    if (result) {
-        Py_DECREF(result);
-        return;
+    // Try first with the stdlib traceback module, but be careful about
+    // traceback module shadowing (issue gh-138170). Use safe import check.
+    if (_PyTraceback_IsSafeToImport()) {
+        PyObject *print_exception_fn = PyImport_ImportModuleAttrString("traceback", "_print_exception_bltin");
+        if (print_exception_fn != NULL && PyCallable_Check(print_exception_fn)) {
+            PyObject* result = PyObject_CallOneArg(print_exception_fn, value);
+            Py_DECREF(print_exception_fn);
+            if (result) {
+                Py_DECREF(result);
+                return;
+            }
+        }
+        else {
+            Py_XDECREF(print_exception_fn);
+        }
     }
 fallback:
 #ifdef Py_DEBUG

Python/traceback.c

@@ -14,6 +14,9 @@
 #include "frameobject.h"          // PyFrame_New()
 
 #include "osdefs.h"               // SEP
+#include <stdio.h>                // fopen, fclose
+#include <limits.h>               // PATH_MAX
+#include <string.h>               // strlen
 #ifdef HAVE_UNISTD_H
 #  include <unistd.h>             // lseek()
 #endif
@@ -69,6 +72,37 @@ class traceback "PyTracebackObject *" "&PyTraceback_Type"
 
 #include "clinic/traceback.c.h"
 
+// Helper function to check if it's safe to import traceback module
+// Returns 0 if a traceback.py file exists in sys.path[0], 1 if safe
+int
+_PyTraceback_IsSafeToImport(void)
+{
+    PyObject *sys_path = PySys_GetObject("path");
+    if (sys_path == NULL || !PyList_Check(sys_path) || PyList_Size(sys_path) == 0) {
+        return 1;  // Default to safe
+    }
+    PyObject *first_path = PyList_GetItem(sys_path, 0);
+    if (first_path == NULL || !PyUnicode_Check(first_path)) {
+        return 1;
+    }
+    const char *path_str = PyUnicode_AsUTF8(first_path);
+    if (path_str == NULL || strlen(path_str) == 0) {
+        return 1;
+    }
+    // Check if traceback.py exists in the first path directory
+    char traceback_path[PATH_MAX];
+    int ret = snprintf(traceback_path, sizeof(traceback_path), "%s/traceback.py", path_str);
+    if (ret <= 0 || ret >= (int)sizeof(traceback_path)) {
+        return 1;  // Path too long or other error, default to safe
+    }
+    FILE *f = fopen(traceback_path, "r");
+    if (f != NULL) {
+        fclose(f);
+        return 0;  // Not safe - traceback.py exists
+    }
+    return 1;  // Safe to import
+}
+
 static PyObject *
 tb_create_raw(PyTracebackObject *next, PyFrameObject *frame, int lasti,
               int lineno)