129 lines
4.1 KiB
Python
129 lines
4.1 KiB
Python
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)
|