import os
import tempfile
import zipfile

from modelaudit.scanners.base import IssueSeverity
from modelaudit.scanners.zip_scanner import ZipScanner


class TestZipScanner:
    """Test the ZIP scanner"""

    def setup_method(self):
        """Set up test fixtures"""
        self.scanner = ZipScanner()

    def test_can_handle_zip_files(self):
        """Test that the scanner correctly identifies ZIP files"""
        # Create a temporary zip file
        with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp:
            with zipfile.ZipFile(tmp.name, "w") as z:
                z.writestr("test.txt", "Hello World")
            tmp_path = tmp.name

        try:
            assert ZipScanner.can_handle(tmp_path) is True
            assert ZipScanner.can_handle("/path/to/file.txt") is False
            assert ZipScanner.can_handle("/path/to/file.pkl") is False
        finally:
            os.unlink(tmp_path)

    def test_symlink_outside_extraction_root(self):
        """Symlinks resolving outside the extraction root should be flagged."""
        with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp:
            with zipfile.ZipFile(tmp.name, "w") as z:
                import stat

                info = zipfile.ZipInfo("link.txt")
                info.create_system = 3
                info.external_attr = (stat.S_IFLNK | 0o777) << 16
                z.writestr(info, "../evil.txt")
            tmp_path = tmp.name

        try:
            result = self.scanner.scan(tmp_path)
            symlink_issues = [i for i in result.issues if "symlink" in i.message.lower()]
            assert any("outside" in i.message.lower() for i in symlink_issues)
        finally:
            os.unlink(tmp_path)

    def test_symlink_to_critical_path(self):
        """Symlinks targeting critical system paths should be flagged."""
        with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp:
            with zipfile.ZipFile(tmp.name, "w") as z:
                import stat

                info = zipfile.ZipInfo("etc_passwd")
                info.create_system = 3
                info.external_attr = (stat.S_IFLNK | 0o777) << 16
                z.writestr(info, "/etc/passwd")
            tmp_path = tmp.name

        try:
            result = self.scanner.scan(tmp_path)
            symlink_issues = [i for i in result.issues if "symlink" in i.message.lower()]
            assert any("critical system" in i.message.lower() for i in symlink_issues)
        finally:
            os.unlink(tmp_path)

    def test_zip_bytes_scanned_single_count(self):
        """Ensure bytes scanned equals the sum of embedded files once."""
        with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp:
            with zipfile.ZipFile(tmp.name, "w") as z:
                import pickle

                data1 = pickle.dumps({"a": 1})
                data2 = pickle.dumps({"b": 2})
                z.writestr("one.pkl", data1)
                z.writestr("two.pkl", data2)
            tmp_path = tmp.name

        try:
            result = self.scanner.scan(tmp_path)
            assert result.success is True
            expected = len(data1) + len(data2)
            assert result.bytes_scanned == expected
        finally:
            os.unlink(tmp_path)

    def test_scan_simple_zip(self):
        """Test scanning a simple ZIP file with text files"""
        with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp:
            with zipfile.ZipFile(tmp.name, "w") as z:
                z.writestr("readme.txt", "This is a readme file")
                z.writestr("data.json", '{"key": "value"}')
            tmp_path = tmp.name

        try:
            result = self.scanner.scan(tmp_path)
            assert result.success is True
            assert result.bytes_scanned > 0
            # May have some debug/info issues about unknown formats
            error_issues = [i for i in result.issues if i.severity == IssueSeverity.CRITICAL]
            assert len(error_issues) == 0
        finally:
            os.unlink(tmp_path)

    def test_scan_zip_with_pickle(self):
        """Test scanning a ZIP file containing a pickle file"""
        with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp:
            with zipfile.ZipFile(tmp.name, "w") as z:
                # Create a simple pickle file
                import pickle

                pickle_data = pickle.dumps({"safe": "data"})
                z.writestr("model.pkl", pickle_data)
            tmp_path = tmp.name

        try:
            result = self.scanner.scan(tmp_path)
            assert result.success is True
            assert result.bytes_scanned > 0
            # The pickle scanner was run on the embedded file
            # Check that we scanned the pickle data
            assert result.bytes_scanned >= len(pickle_data)
        finally:
            os.unlink(tmp_path)

    def test_scan_nested_zip(self):
        """Test scanning nested ZIP files"""
        # Create inner zip
        with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as inner_tmp:
            with zipfile.ZipFile(inner_tmp.name, "w") as inner_z:
                inner_z.writestr("inner.txt", "Inner file content")
            inner_path = inner_tmp.name

        try:
            # Create outer zip containing inner zip
            with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as outer_tmp:
                with zipfile.ZipFile(outer_tmp.name, "w") as outer_z:
                    outer_z.write(inner_path, "nested.zip")
                outer_path = outer_tmp.name

            result = self.scanner.scan(outer_path)
            assert result.success is True
            # Should have scanned the nested content
            assert (
                any("nested.zip" in str(issue.location) for issue in result.issues if hasattr(issue, "location"))
                or result.bytes_scanned > 0
            )
        finally:
            os.unlink(inner_path)
            os.unlink(outer_path)

    def test_directory_traversal_detection(self):
        """Test detection of directory traversal attempts in ZIP files"""
        with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp:
            with zipfile.ZipFile(tmp.name, "w") as z:
                # Create entries with directory traversal attempts
                z.writestr("../../../etc/passwd", "malicious content")
                z.writestr("/etc/passwd", "malicious content")
                z.writestr("safe.txt", "safe content")
            tmp_path = tmp.name

        try:
            result = self.scanner.scan(tmp_path)
            assert result.success is True

            # Should have detected directory traversal attempts
            traversal_issues = [
                i
                for i in result.issues
                if "path traversal" in i.message.lower() or "directory traversal" in i.message.lower()
            ]
            assert len(traversal_issues) >= 2

            # Check severity
            for issue in traversal_issues:
                assert issue.severity == IssueSeverity.CRITICAL
        finally:
            os.unlink(tmp_path)

    def test_windows_traversal_detection(self):
        """Ensure Windows-style path traversal is caught"""
        with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp:
            with zipfile.ZipFile(tmp.name, "w") as z:
                z.writestr("..\\evil.txt", "malicious")
                z.writestr("safe.txt", "ok")
            tmp_path = tmp.name

        try:
            result = self.scanner.scan(tmp_path)
            traversal_issues = [i for i in result.issues if "path traversal" in i.message.lower()]
            assert len(traversal_issues) >= 1
            for issue in traversal_issues:
                assert issue.severity == IssueSeverity.CRITICAL
        finally:
            os.unlink(tmp_path)

    def test_zip_bomb_detection(self):
        """Test detection of potential zip bombs (high compression ratio)"""
        with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp:
            with zipfile.ZipFile(tmp.name, "w", compression=zipfile.ZIP_DEFLATED) as z:
                # Create a highly compressible file (potential zip bomb indicator)
                # Keep highly compressible but smaller to speed CI
                large_content = "A" * 300000  # 300KB of repeated 'A's
                z.writestr("suspicious.txt", large_content)
            tmp_path = tmp.name

        try:
            result = self.scanner.scan(tmp_path)
            assert result.success is True

            # Should detect high compression ratio
            compression_issues = [i for i in result.issues if "compression ratio" in i.message.lower()]
            assert len(compression_issues) >= 1
        finally:
            os.unlink(tmp_path)

    def test_max_depth_limit(self):
        """Test that maximum nesting depth is enforced"""
        # Create deeply nested zips
        current_path = None
        paths_to_delete = []

        try:
            # Create 10 levels of nested zips
            for i in range(10):
                with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp:
                    with zipfile.ZipFile(tmp.name, "w") as z:
                        if current_path:
                            z.write(current_path, f"level{i}.zip")
                        else:
                            z.writestr("deepest.txt", "Deep content")
                    paths_to_delete.append(tmp.name)
                    current_path = tmp.name

            # Scan the outermost zip
            assert current_path is not None  # Should be set by the loop above
            scanner = ZipScanner(config={"max_zip_depth": 3})
            result = scanner.scan(current_path)

            assert result.success is True
            # Should have a warning about max depth
            depth_issues = [i for i in result.issues if "depth" in i.message.lower()]
            assert len(depth_issues) >= 1
        finally:
            for path in paths_to_delete:
                if os.path.exists(path):
                    os.unlink(path)

    def test_scan_zip_with_dangerous_pickle(self):
        """Test scanning a ZIP file containing a dangerous pickle"""
        with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp:
            with zipfile.ZipFile(tmp.name, "w") as z:
                # Create a pickle with suspicious content
                import os as os_module
                import pickle

                class DangerousClass:
                    def __reduce__(self):
                        return (os_module.system, ("echo pwned",))

                dangerous_obj = DangerousClass()
                pickle_data = pickle.dumps(dangerous_obj)
                z.writestr("dangerous.pkl", pickle_data)
            tmp_path = tmp.name

        try:
            result = self.scanner.scan(tmp_path)
            # The scan should complete even if there are errors in the pickle scanner
            assert result.success is True

            # Check that we at least tried to scan the pickle
            assert result.bytes_scanned > 0

            # May have error issues due to the bug in pickle scanner with string_stack
            # or it may detect the dangerous content
            # Either way, it should have scanned the file
        finally:
            os.unlink(tmp_path)

    def test_scan_nonexistent_file(self):
        """Test scanning a file that doesn't exist"""
        result = self.scanner.scan("/nonexistent/file.zip")
        assert result.success is False
        assert len(result.issues) > 0
        assert any("does not exist" in issue.message for issue in result.issues)

    def test_scan_invalid_zip(self):
        """Test scanning a file that's not a valid ZIP"""
        with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp:
            tmp.write(b"This is not a zip file")
            tmp_path = tmp.name

        try:
            result = self.scanner.scan(tmp_path)
            assert result.success is False
            assert len(result.issues) > 0
            assert any("not a valid zip" in issue.message.lower() for issue in result.issues)
        finally:
            os.unlink(tmp_path)
