AAP-78319 Fix path traversal in awxkit YAML (#16608)

* Fix path traversal in awxkit YAML

* coderabbitai feedback

awxkit/awxkit/exceptions.py

@@ -89,3 +89,7 @@ class IsMigrating(Common):
 
 class ImportExportError(Exception):
     pass
+
+
+class PathTraversalError(Exception):
+    pass

awxkit/awxkit/yaml_file.py

@@ -3,6 +3,8 @@ import yaml
 import glob
 import logging
 
+from awxkit.exceptions import PathTraversalError
+
 log = logging.getLogger(__name__)
 
 
@@ -12,7 +14,11 @@ file_path_cache = {}
 
 class Loader(yaml.SafeLoader):
     def __init__(self, stream):
-        self._root = os.path.split(stream.name)[0]
+        if stream.name == '<stdin>':
+            # !include is unsupported when reading from stdin; there's no base directory to resolve relative paths against
+            self._root = ''
+        else:
+            self._root = os.path.dirname(os.path.realpath(stream.name))
         super(Loader, self).__init__(stream)
         Loader.add_constructor('!include', Loader.include)
         Loader.add_constructor('!import', Loader.include)
@@ -37,8 +43,25 @@ class Loader(yaml.SafeLoader):
             log.error("unrecognised node type in !include statement")
             raise yaml.constructor.ConstructorError
 
+    def _validate_path(self, filename, resolved_path):
+        real_path = os.path.realpath(resolved_path)
+        root_prefix = self._root + os.sep
+        if not (real_path.startswith(root_prefix) or real_path == self._root):
+            raise PathTraversalError("YAML !include path '{}' is not allowed".format(filename))
+
     def extractFile(self, filename):
+        if not self._root:
+            raise PathTraversalError("YAML !include directive is not supported when reading from stdin")
+
+        if os.path.isabs(filename):
+            raise PathTraversalError("YAML !include path '{}' is not allowed".format(filename))
+
         file_pattern = os.path.join(self._root, filename)
+
+        normalized = os.path.normpath(file_pattern)
+        root_prefix = self._root + os.sep
+        if not (normalized.startswith(root_prefix) or normalized == self._root):
+            raise PathTraversalError("YAML !include path '{}' is not allowed".format(filename))
         log.debug('Will attempt to extract schema from: {0}'.format(file_pattern))
         if file_pattern in file_pattern_cache:
             log.debug('File pattern cache hit: {0}'.format(file_pattern))
@@ -47,6 +70,7 @@ class Loader(yaml.SafeLoader):
         data = dict()
         for file_path in glob.glob(file_pattern):
             file_path = os.path.abspath(file_path)
+            self._validate_path(filename, file_path)
             if file_path in file_path_cache:
                 log.debug('Schema cache hit: {0}'.format(file_path))
                 path_data = file_path_cache[file_path]

awxkit/test/test_yaml_file.py

@@ -0,0 +1,148 @@
+import io
+import os
+
+import pytest
+import yaml
+
+from awxkit.exceptions import PathTraversalError
+from awxkit.yaml_file import Loader, file_path_cache, file_pattern_cache, load_file
+
+
+@pytest.fixture(autouse=True)
+def clear_yaml_caches():
+    file_pattern_cache.clear()
+    file_path_cache.clear()
+    yield
+    file_pattern_cache.clear()
+    file_path_cache.clear()
+
+
+@pytest.fixture
+def yaml_tree(tmp_path):
+    project = tmp_path / "project"
+    project.mkdir()
+    sub = project / "sub"
+    sub.mkdir()
+    outside = tmp_path / "outside"
+    outside.mkdir()
+
+    (project / "main.yaml").write_text("key: value\n")
+    (sub / "included.yaml").write_text("included_key: included_value\n")
+    (sub / "nested.yaml").write_text("nested_key: nested_value\n")
+    (outside / "secret.yaml").write_text("secret: exposed\n")
+
+    return project
+
+
+def load_yaml_with_root(content, root_dir):
+    stream = io.StringIO(content)
+    stream.name = os.path.join(root_dir, "test_input.yaml")
+    return yaml.load(stream, Loader=Loader)
+
+
+def load_yaml_as_stdin(content):
+    stream = io.StringIO(content)
+    stream.name = "<stdin>"
+    return yaml.load(stream, Loader=Loader)
+
+
+class TestIncludeLegitimate:
+    def test_relative_same_directory(self, yaml_tree):
+        result = load_yaml_with_root("data: !include sub/included.yaml", str(yaml_tree))
+        assert result["data"]["included_key"] == "included_value"
+
+    def test_glob_within_tree(self, yaml_tree):
+        result = load_yaml_with_root("data: !include sub/*.yaml", str(yaml_tree))
+        assert "included_key" in result["data"]
+        assert "nested_key" in result["data"]
+
+    def test_sequence_node(self, yaml_tree):
+        content = "data: !include\n  - sub/included.yaml\n  - sub/nested.yaml\n"
+        result = load_yaml_with_root(content, str(yaml_tree))
+        assert "included_key" in result["data"]
+        assert "nested_key" in result["data"]
+
+    def test_mapping_node(self, yaml_tree):
+        content = "data: !include\n  included_key: sub/included.yaml\n  nested_key: sub/nested.yaml\n"
+        result = load_yaml_with_root(content, str(yaml_tree))
+        assert result["data"]["included_key"] == "included_value"
+        assert result["data"]["nested_key"] == "nested_value"
+
+    def test_load_file_with_include(self, yaml_tree):
+        main_yaml = yaml_tree / "main_with_include.yaml"
+        main_yaml.write_text("imported: !include sub/included.yaml\n")
+        result = load_file(str(main_yaml))
+        assert result["imported"]["included_key"] == "included_value"
+
+    def test_relative_filename_with_no_directory_component(self, yaml_tree, monkeypatch):
+        # A stream.name with no directory component (e.g. `open("main.yaml")` from cwd)
+        # must still resolve _root to cwd, not be mistaken for stdin.
+        monkeypatch.chdir(yaml_tree)
+        stream = io.StringIO("data: !include sub/included.yaml")
+        stream.name = "main.yaml"
+        result = yaml.load(stream, Loader=Loader)
+        assert result["data"]["included_key"] == "included_value"
+
+
+class TestIncludePathTraversal:
+    def test_relative_traversal(self, yaml_tree):
+        with pytest.raises(PathTraversalError):
+            load_yaml_with_root("data: !include ../outside/secret.yaml", str(yaml_tree))
+
+    def test_dotdot_deeper(self, yaml_tree):
+        with pytest.raises(PathTraversalError):
+            load_yaml_with_root("data: !include ../../etc/passwd", str(yaml_tree))
+
+    def test_absolute_path(self, yaml_tree):
+        with pytest.raises(PathTraversalError):
+            load_yaml_with_root("data: !include /etc/passwd", str(yaml_tree))
+
+    def test_glob_traversal(self, yaml_tree):
+        with pytest.raises(PathTraversalError):
+            load_yaml_with_root("data: !include ../outside/*.yaml", str(yaml_tree))
+
+    def test_stdin_rejected(self):
+        with pytest.raises(PathTraversalError, match="stdin"):
+            load_yaml_as_stdin("data: !include some_file.yaml")
+
+    def test_symlink_escape(self, yaml_tree, tmp_path):
+        outside_file = tmp_path / "outside" / "secret.yaml"
+        symlink = yaml_tree / "sub" / "link_to_outside.yaml"
+        symlink.symlink_to(outside_file)
+
+        with pytest.raises(PathTraversalError):
+            load_yaml_with_root("data: !include sub/link_to_outside.yaml", str(yaml_tree))
+
+    def test_traversal_in_sequence(self, yaml_tree):
+        content = "data: !include\n  - sub/included.yaml\n  - ../outside/secret.yaml\n"
+        with pytest.raises(PathTraversalError):
+            load_yaml_with_root(content, str(yaml_tree))
+
+    def test_error_message_says_not_allowed(self, yaml_tree):
+        with pytest.raises(PathTraversalError, match="not allowed"):
+            load_yaml_with_root("data: !include ../outside/secret.yaml", str(yaml_tree))
+
+    def test_traversal_in_mapping_value(self, yaml_tree):
+        content = "data: !include\n  included_key: ../outside/secret.yaml\n"
+        with pytest.raises(PathTraversalError):
+            load_yaml_with_root(content, str(yaml_tree))
+
+    def test_symlink_directory_escape(self, yaml_tree, tmp_path):
+        outside_dir = tmp_path / "outside"
+        link_dir = yaml_tree / "sub" / "link_to_outside"
+        link_dir.symlink_to(outside_dir)
+
+        with pytest.raises(PathTraversalError):
+            load_yaml_with_root("data: !include sub/link_to_outside/secret.yaml", str(yaml_tree))
+
+    def test_load_file_with_traversal(self, yaml_tree):
+        malicious = yaml_tree / "malicious.yaml"
+        malicious.write_text("exploit: !include ../outside/secret.yaml\n")
+        with pytest.raises(PathTraversalError):
+            load_file(str(malicious))
+
+    def test_error_message_absolute_path(self):
+        stream = io.StringIO("data: !include /etc/passwd")
+        stream.name = "/some/dir/test.yaml"
+        with pytest.raises(PathTraversalError, match="not allowed"):
+            yaml.load(stream, Loader=Loader)