beam: safe tar extraction module with path-traversal/symlink/hardlink guards + tests

wizard-work/beam/safe_extract.py

@@ -0,0 +1,171 @@
+"""Safe tar extraction with path-traversal, symlink and hardlink guards.
+
+Public API:
+    safe_extract(archive, dest, *, max_members=10000, max_file_bytes=None)
+    SafeExtractError
+
+Raises SafeExtractError on any unsafe member. Never writes outside `dest`.
+"""
+import os
+import tarfile
+
+__all__ = ["safe_extract", "SafeExtractError"]
+
+
+class SafeExtractError(Exception):
+    pass
+
+
+def _norm_rel(name):
+    if not isinstance(name, str) or name == "":
+        raise SafeExtractError("empty member name")
+    # tar names may use backslashes on windows-ish archives; normalize
+    name = name.replace("\\", "/")
+    if name.startswith("/"):
+        raise SafeExtractError("absolute member path: %r" % name)
+    parts = name.split("/")
+    # drop a single leading "./"
+    if parts and parts[0] == ".":
+        parts = parts[1:]
+    if not parts or all(p == "" for p in parts):
+        raise SafeExtractError("empty member path: %r" % name)
+    if any(p in ("", ".", "..") for p in parts):
+        raise SafeExtractError("traversal/empty component in member path: %r" % name)
+    return os.path.join(*parts)
+
+
+def _ensure_within(dest, target):
+    dest = os.path.abspath(dest)
+    target = os.path.abspath(target)
+    if target == dest:
+        return target
+    if not (target.startswith(dest + os.sep)):
+        raise SafeExtractError("member escapes destination: %r" % target)
+    return target
+
+
+def _safe_link_target(dest, member_dir, linkname):
+    if not isinstance(linkname, str) or linkname == "":
+        raise SafeExtractError("empty link target")
+    linkname = linkname.replace("\\", "/")
+    if linkname.startswith("/"):
+        raise SafeExtractError("absolute link target: %r" % linkname)
+    # resolve relative to member's directory
+    joined = os.path.normpath(os.path.join(member_dir, linkname))
+    # member_dir is already absolute inside dest
+    if joined == os.path.abspath(dest):
+        return joined
+    if not joined.startswith(os.path.abspath(dest) + os.sep):
+        raise SafeExtractError("link target escapes destination: %r" % linkname)
+    return joined
+
+
+def safe_extract(archive, dest, *, max_members=10000, max_file_bytes=None):
+    dest = os.path.abspath(dest)
+    os.makedirs(dest, exist_ok=True)
+    # refuse if dest itself is a symlink to elsewhere? We require dest real.
+    if os.path.islink(dest):
+        raise SafeExtractError("destination is a symlink")
+
+    with tarfile.open(archive, "r:*") as tf:
+        members = tf.getmembers()
+        if len(members) > max_members:
+            raise SafeExtractError("too many members: %d > %d" % (len(members), max_members))
+
+        # First pass: validate all members before writing anything.
+        planned = []
+        for m in members:
+            rel = _norm_rel(m.name)
+            target = _ensure_within(dest, os.path.join(dest, rel))
+            if m.issym() or m.islnk():
+                member_dir = os.path.dirname(target)
+                _safe_link_target(dest, member_dir, m.linkname)
+            if max_file_bytes is not None and m.isfile():
+                if m.size > max_file_bytes:
+                    raise SafeExtractError("file too large: %s (%d bytes)" % (m.name, m.size))
+            planned.append((m, target))
+
+        # Second pass: extract, never following symlinks for parents.
+        for m, target in planned:
+            # ensure parent exists and is not a symlink escaping dest
+            parent = os.path.dirname(target)
+            # walk/create parents manually, refusing symlink parents
+            cur = dest
+            for part in os.path.relpath(parent, dest).split(os.sep):
+                if part in ("", "."):
+                    continue
+                cur = os.path.join(cur, part)
+                if os.path.lexists(cur):
+                    if os.path.islink(cur):
+                        raise SafeExtractError("refusing to traverse symlinked parent: %r" % cur)
+                else:
+                    os.mkdir(cur)
+            _ensure_within(dest, cur)
+
+            if m.isdir():
+                if os.path.lexists(target):
+                    if os.path.islink(target):
+                        raise SafeExtractError("refusing to replace symlink with dir: %r" % target)
+                else:
+                    os.mkdir(target)
+            elif m.issym():
+                # remove existing then create link
+                if os.path.lexists(target):
+                    if os.path.isdir(target) and not os.path.islink(target):
+                        raise SafeExtractError("refusing to overwrite directory with symlink: %r" % target)
+                    os.remove(target) if not os.path.isdir(target) or os.path.islink(target) else None
+                if os.path.lexists(target):
+                    # directory case handled above; best effort remove
+                    try:
+                        os.rmdir(target)
+                    except OSError:
+                        raise SafeExtractError("cannot replace existing path with symlink: %r" % target)
+                linktarget = m.linkname.replace("\\", "/")
+                # store relative/absolute as validated; use validated resolved? keep original relative
+                os.symlink(linktarget, target)
+            elif m.islnk():
+                # hardlink: resolve source within dest
+                src_rel = _norm_rel(m.linkname) if not os.path.isabs(m.linkname) else None
+                # hardlink linkname is a member name reference
+                if os.path.isabs(m.linkname):
+                    raise SafeExtractError("absolute hardlink target: %r" % m.linkname)
+                src = _ensure_within(dest, os.path.join(dest, _norm_rel(m.linkname)))
+                if os.path.islink(src):
+                    raise SafeExtractError("hardlink source is symlink: %r" % src)
+                # copy file content instead of hardlink to avoid cross-escape
+                if os.path.isdir(src):
+                    raise SafeExtractError("hardlink to directory not allowed")
+                # ensure target parent done; copy
+                if os.path.lexists(target):
+                    if os.path.isdir(target) and not os.path.islink(target):
+                        raise SafeExtractError("refusing to overwrite dir with file: %r" % target)
+                    os.remove(target)
+                import shutil as _sh
+                _sh.copyfile(src, target)
+            elif m.isfile():
+                if os.path.lexists(target) and os.path.isdir(target) and not os.path.islink(target):
+                    raise SafeExtractError("refusing to overwrite directory with file: %r" % target)
+                if os.path.lexists(target):
+                    os.remove(target)
+                srcf = tf.extractfile(m)
+                if srcf is None:
+                    raise SafeExtractError("cannot read member: %r" % m.name)
+                with open(target, "wb") as out:
+                    remaining = m.size
+                    chunk = 65536
+                    while remaining > 0:
+                        data = srcf.read(min(chunk, remaining))
+                        if not data:
+                            break
+                        out.write(data)
+                        remaining -= len(data)
+                # mode bits (best effort)
+                try:
+                    os.chmod(target, m.mode & 0o7777)
+                except OSError:
+                    pass
+            else:
+                # devices, fifos, etc. skip for safety
+                raise SafeExtractError("unsupported member type: %r" % m.name)
+
+    return dest

wizard-work/beam/test_safe_extract.py

@@ -0,0 +1,128 @@
+import os, tarfile, tempfile, unittest, shutil
+import sys
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from safe_extract import safe_extract, SafeExtractError
+
+
+def make_tar(members_fn, path):
+    with tarfile.open(path, "w:gz") as tf:
+        members_fn(tf)
+
+
+class TestSafeExtract(unittest.TestCase):
+    def setUp(self):
+        self.tmp = tempfile.mkdtemp()
+        self.dest = os.path.join(self.tmp, "out")
+        os.makedirs(self.dest)
+
+    def tearDown(self):
+        shutil.rmtree(self.tmp, ignore_errors=True)
+
+    def test_basic_file(self):
+        arc = os.path.join(self.tmp, "a.tar.gz")
+        def fn(tf):
+            import io
+            data = b"hello"
+            ti = tarfile.TarInfo("hello.txt")
+            ti.size = len(data)
+            tf.addfile(ti, io.BytesIO(data))
+        make_tar(fn, arc)
+        safe_extract(arc, self.dest)
+        self.assertEqual(open(os.path.join(self.dest, "hello.txt"), "rb").read(), b"hello")
+
+    def test_traversal_rejected(self):
+        arc = os.path.join(self.tmp, "b.tar.gz")
+        def fn(tf):
+            import io
+            data = b"x"
+            ti = tarfile.TarInfo("../evil.txt")
+            ti.size = len(data)
+            tf.addfile(ti, io.BytesIO(data))
+        make_tar(fn, arc)
+        with self.assertRaises(SafeExtractError):
+            safe_extract(arc, self.dest)
+        self.assertFalse(os.path.exists(os.path.join(self.tmp, "evil.txt")))
+
+    def test_absolute_rejected(self):
+        arc = os.path.join(self.tmp, "c.tar.gz")
+        def fn(tf):
+            import io
+            data = b"x"
+            ti = tarfile.TarInfo("/etc/evil.txt")
+            ti.size = len(data)
+            tf.addfile(ti, io.BytesIO(data))
+        make_tar(fn, arc)
+        with self.assertRaises(SafeExtractError):
+            safe_extract(arc, self.dest)
+
+    def test_symlink_escape_rejected(self):
+        arc = os.path.join(self.tmp, "d.tar.gz")
+        def fn(tf):
+            import io
+            ti = tarfile.TarInfo("link")
+            ti.type = tarfile.SYMTYPE
+            ti.linkname = "../../etc/passwd"
+            tf.addfile(ti)
+        make_tar(fn, arc)
+        with self.assertRaises(SafeExtractError):
+            safe_extract(arc, self.dest)
+
+    def test_internal_symlink_ok(self):
+        arc = os.path.join(self.tmp, "e.tar.gz")
+        def fn(tf):
+            import io
+            data = b"hi"
+            ti = tarfile.TarInfo("a.txt")
+            ti.size = len(data)
+            tf.addfile(ti, io.BytesIO(data))
+            ti2 = tarfile.TarInfo("l")
+            ti2.type = tarfile.SYMTYPE
+            ti2.linkname = "a.txt"
+            tf.addfile(ti2)
+        make_tar(fn, arc)
+        safe_extract(arc, self.dest)
+        self.assertTrue(os.path.islink(os.path.join(self.dest, "l")))
+
+    def test_hardlink_escape_rejected(self):
+        arc = os.path.join(self.tmp, "f.tar.gz")
+        def fn(tf):
+            import io
+            data = b"hi"
+            ti = tarfile.TarInfo("a.txt")
+            ti.size = len(data)
+            tf.addfile(ti, io.BytesIO(data))
+            ti2 = tarfile.TarInfo("h")
+            ti2.type = tarfile.LNKTYPE
+            ti2.linkname = "../outside.txt"
+            tf.addfile(ti2)
+        make_tar(fn, arc)
+        with self.assertRaises(SafeExtractError):
+            safe_extract(arc, self.dest)
+
+    def test_size_limit(self):
+        arc = os.path.join(self.tmp, "g.tar.gz")
+        def fn(tf):
+            import io
+            data = b"x" * 100
+            ti = tarfile.TarInfo("big.txt")
+            ti.size = len(data)
+            tf.addfile(ti, io.BytesIO(data))
+        make_tar(fn, arc)
+        with self.assertRaises(SafeExtractError):
+            safe_extract(arc, self.dest, max_file_bytes=10)
+
+    def test_device_rejected(self):
+        arc = os.path.join(self.tmp, "h.tar.gz")
+        def fn(tf):
+            ti = tarfile.TarInfo("dev")
+            ti.type = tarfile.CHRTYPE
+            ti.devmajor = 1
+            ti.devminor = 3
+            tf.addfile(ti)
+        make_tar(fn, arc)
+        with self.assertRaises(SafeExtractError):
+            safe_extract(arc, self.dest)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)