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)