omnigent-ai--omnigent/tests/spec/test_tar_utils.py

"""Tests for omnigent.spec.tar_utils."""

from __future__ import annotations

import io
import tarfile
from pathlib import Path

import pytest
import yaml

from omnigent.spec.tar_utils import ExtractionError, extract_safe


def _create_tar(tmp_path: Path, members: dict[str, bytes | str]) -> Path:
    """Create a tar.gz at tmp_path/bundle.tar.gz with given members."""
    tar_path = tmp_path / "bundle.tar.gz"
    with tarfile.open(tar_path, "w:gz") as tf:
        for name, content in members.items():
            data = content.encode() if isinstance(content, str) else content
            info = tarfile.TarInfo(name=name)
            info.size = len(data)
            tf.addfile(info, io.BytesIO(data))
    return tar_path


@pytest.fixture()
def dest(tmp_path: Path) -> Path:
    return tmp_path / "extracted"


def test_extract_valid_tarball(tmp_path: Path, dest: Path) -> None:
    config = yaml.dump({"spec_version": 1, "name": "test"})
    tar_path = _create_tar(tmp_path, {"config.yaml": config})
    result = extract_safe(tar_path, dest)
    assert result == dest
    assert (dest / "config.yaml").exists()
    assert yaml.safe_load((dest / "config.yaml").read_text())["name"] == "test"


def test_extract_nested_files(tmp_path: Path, dest: Path) -> None:
    tar_path = _create_tar(
        tmp_path,
        {
            "config.yaml": "spec_version: 1",
            "skills/search/SKILL.md": "---\nname: search\n---\ncontent",
        },
    )
    extract_safe(tar_path, dest)
    assert (dest / "skills" / "search" / "SKILL.md").exists()


def test_extract_rejects_path_traversal(tmp_path: Path, dest: Path) -> None:
    tar_path = _create_tar(tmp_path, {"../escape.txt": "evil"})
    with pytest.raises(ExtractionError, match="path traversal"):
        extract_safe(tar_path, dest)


def test_extract_rejects_absolute_path(tmp_path: Path, dest: Path) -> None:
    tar_path = _create_tar(tmp_path, {"/etc/passwd": "evil"})
    with pytest.raises(ExtractionError, match="absolute path"):
        extract_safe(tar_path, dest)


def test_extract_rejects_symlink(tmp_path: Path, dest: Path) -> None:
    tar_path = tmp_path / "symlink.tar.gz"
    with tarfile.open(tar_path, "w:gz") as tf:
        info = tarfile.TarInfo(name="evil-link")
        info.type = tarfile.SYMTYPE
        info.linkname = "/etc/passwd"
        tf.addfile(info)
    with pytest.raises(ExtractionError, match="link"):
        extract_safe(tar_path, dest)


def test_extract_rejects_hardlink(tmp_path: Path, dest: Path) -> None:
    tar_path = tmp_path / "hardlink.tar.gz"
    with tarfile.open(tar_path, "w:gz") as tf:
        info = tarfile.TarInfo(name="evil-link")
        info.type = tarfile.LNKTYPE
        info.linkname = "/etc/passwd"
        tf.addfile(info)
    with pytest.raises(ExtractionError, match="link"):
        extract_safe(tar_path, dest)


@pytest.mark.parametrize(
    "member_type,type_label",
    [
        (tarfile.FIFOTYPE, "FIFO"),
        (tarfile.CHRTYPE, "character device"),
        (tarfile.BLKTYPE, "block device"),
    ],
)
def test_extract_rejects_special_file_types(
    tmp_path: Path, dest: Path, member_type: bytes, type_label: str
) -> None:
    """
    Special tar members (FIFOs, char/block devices) must be
    rejected by the allow-list, and must NOT be created on disk.

    The original deny-list only blocked links, so a FIFO named
    ``config.yaml`` was materialized on disk; the spec loader then
    hung forever on ``read_text()`` waiting for a writer, exhausting
    the AP-server worker pool.
    """
    tar_path = tmp_path / f"special-{type_label.replace(' ', '_')}.tar.gz"
    with tarfile.open(tar_path, "w:gz") as tf:
        info = tarfile.TarInfo(name="config.yaml")
        info.type = member_type
        tf.addfile(info)

    with pytest.raises(ExtractionError, match="unsupported entry type"):
        extract_safe(tar_path, dest)

    # The whole point of the fix: the special node is never created,
    # so there is no FIFO for a later read_text() to block on. If this
    # fails, the node was written before the check rejected it and the
    # DoS vector is still open.
    assert not (dest / "config.yaml").exists()


def test_extract_rejects_special_member_alongside_valid_files(tmp_path: Path, dest: Path) -> None:
    """
    A bundle that mixes a valid ``config.yaml`` with a FIFO entry is
    rejected wholesale — extraction must not partially succeed and
    leave the FIFO behind.
    """
    tar_path = tmp_path / "mixed.tar.gz"
    with tarfile.open(tar_path, "w:gz") as tf:
        good = b"spec_version: 1\nname: test\n"
        good_info = tarfile.TarInfo(name="config.yaml")
        good_info.size = len(good)
        tf.addfile(good_info, io.BytesIO(good))

        fifo_info = tarfile.TarInfo(name="pipe")
        fifo_info.type = tarfile.FIFOTYPE
        tf.addfile(fifo_info)

    with pytest.raises(ExtractionError, match="unsupported entry type"):
        extract_safe(tar_path, dest)

    # No FIFO node materialized regardless of member ordering.
    assert not (dest / "pipe").exists()


def test_extract_rejects_size_bomb(tmp_path: Path, dest: Path) -> None:
    tar_path = tmp_path / "bomb.tar.gz"
    with tarfile.open(tar_path, "w:gz") as tf:
        data = b"x" * 1024
        info = tarfile.TarInfo(name="big.bin")
        info.size = len(data)
        tf.addfile(info, io.BytesIO(data))
    with pytest.raises(ExtractionError, match="max extracted size"):
        extract_safe(tar_path, dest, max_bytes=512)


def test_extract_rejects_entry_bomb(tmp_path: Path, dest: Path) -> None:
    tar_path = tmp_path / "entries.tar.gz"
    with tarfile.open(tar_path, "w:gz") as tf:
        for i in range(10):
            info = tarfile.TarInfo(name=f"file_{i}.txt")
            info.size = 1
            tf.addfile(info, io.BytesIO(b"x"))
    with pytest.raises(ExtractionError, match="max entry count"):
        extract_safe(tar_path, dest, max_entries=5)


def test_extract_missing_tarball(tmp_path: Path, dest: Path) -> None:
    with pytest.raises(FileNotFoundError, match="tarball not found"):
        extract_safe(tmp_path / "nonexistent.tar.gz", dest)


def test_extract_creates_dest_directory(tmp_path: Path) -> None:
    dest = tmp_path / "deep" / "nested" / "dir"
    tar_path = _create_tar(tmp_path, {"config.yaml": "spec_version: 1"})
    extract_safe(tar_path, dest)
    assert dest.is_dir()
    assert (dest / "config.yaml").exists()


def test_extract_from_bytes(tmp_path: Path, dest: Path) -> None:
    config = yaml.dump({"spec_version": 1, "name": "bytes-agent"})
    tar_path = _create_tar(tmp_path, {"config.yaml": config})
    bundle_bytes = tar_path.read_bytes()

    result = extract_safe(bundle_bytes, dest)
    assert result == dest
    assert (dest / "config.yaml").exists()
    assert yaml.safe_load((dest / "config.yaml").read_text())["name"] == "bytes-agent"


def test_extract_from_bytes_rejects_traversal(dest: Path) -> None:
    buf = io.BytesIO()
    with tarfile.open(fileobj=buf, mode="w:gz") as tf:
        data = b"evil"
        info = tarfile.TarInfo(name="../escape.txt")
        info.size = len(data)
        tf.addfile(info, io.BytesIO(data))

    with pytest.raises(ExtractionError, match="path traversal"):
        extract_safe(buf.getvalue(), dest)