CTF-tool/tests/test_forensics.py

from pathlib import Path
import toml
import os
import stat
import sys
from click.testing import CliRunner
from ctf.utils import load_config, active_categories
from ctf.forensics import get_metadata
from ctf.cli_forensics import forensics_group, inspect

# Define the persistent test environment path
TEST_ENV = Path("tests/env")

# =====================================================================
# 1. Test Strategy Summary
# =====================================================================
# Core scenarios tested:
#
# A. Happy Path:
#    - Verifying metadata attributes (permissions, ownership, size, 
#      hard links, inode, device) map exactly to filesystem ground truth.
#
# B. Boundary Conditions:
#    - Testing empty files (0 bytes) for correct apparent size and allocation.
#    - Testing file permissions changes (e.g. read-only, executable) and
#      making sure the parsed octal/symbolic codes match.
#    - Testing multiple hard links (count > 1).
#
# C. Edge Cases & OS Portability:
#    - Extended attributes (xattr) support, handling platforms where xattr
#      is not present or supported on the active file system.
#    - Graceful system identity fallbacks for UID/GID names when passwd/group
#      databases are unavailable or running on non-Unix systems.
#
# D. Error Handling:
#    - Accessing non-existent paths (CLI validation errors).
# =====================================================================

# =====================================================================
# 2. Mocking/Setup Requirements
# =====================================================================
# - Filesystem sandbox: Uses `tests/env/` to dynamically create files
#   with specific mode bits, content, and links.
# - No external network/API mocking is required.
# - OS-specific conditional blocks handle platforms (e.g., Windows) 
#   where POSIX ownership resolution or xattr is not native.
# =====================================================================

def test_load_config_static():
    """Verifies load_config using the persistent test file."""
    config_file = TEST_ENV / "config.toml"
    test_data = {
        "Competition": {"name": "PersistentComp"},
        "Enviroment": {"data_dir": str(TEST_ENV.absolute())}
    }
    config_file.parent.mkdir(parents=True, exist_ok=True)
    with open(config_file, "w") as f:
        toml.dump(test_data, f)
    
    loaded_data = load_config(str(config_file))
    assert loaded_data["Competition"]["name"] == "PersistentComp"

def test_active_categories_static():
    """Verifies active_categories using the pre-created 'comp1' folder."""
    comp_dir = TEST_ENV / "comp1"
    comp_dir.mkdir(parents=True, exist_ok=True)
    (comp_dir / "web").mkdir(exist_ok=True)
    (comp_dir / "pwn").mkdir(exist_ok=True)
    
    cats = active_categories(comp_dir)
    cat_names = [c.name for c in cats]
    
    assert "web" in cat_names
    assert "pwn" in cat_names


# --- Task-Specific Universal Metadata Tests ---

def test_inspect_permissions():
    """
    Task 1: POSIX Permissions.
    Verifies that the parsed octal and symbolic representation of permissions
    matches Python's native os.stat mode decoding.
    """
    test_file = TEST_ENV / "perm_test.bin"
    with open(test_file, "wb") as f:
        f.write(b"data")
    
    # Set permissions explicitly to 0o755 (rwxr-xr-x)
    test_file.chmod(0o755)
    
    meta = get_metadata(test_file)
    
    # Assert correct octal format
    assert meta.permissions_octal == "0o755"
    # Assert correct symbolic format (standard file type prefix '-')
    assert meta.permissions_symbolic == "-rwxr-xr-x"
    
    # Change permissions to 0o644 (rw-r--r--)
    test_file.chmod(0o644)
    meta_new = get_metadata(test_file)
    assert meta_new.permissions_octal == "0o644"
    assert meta_new.permissions_symbolic == "-rw-r--r--"


def test_inspect_ownership():
    """
    Task 2: Ownership Identity.
    Verifies that numeric UID/GID and resolved username/groupname match.
    """
    test_file = TEST_ENV / "owner_test.bin"
    with open(test_file, "wb") as f:
        f.write(b"data")
        
    stat_info = test_file.stat()
    meta = get_metadata(test_file)
    
    assert meta.owner_uid == stat_info.st_uid
    assert meta.owner_gid == stat_info.st_gid
    
    # On Unix, verify that user and group strings are resolved
    if sys.platform != "win32":
        import pwd
        import grp
        expected_user = pwd.getpwuid(stat_info.st_uid).pw_name
        expected_group = grp.getgrgid(stat_info.st_gid).gr_name
        assert meta.owner_username == expected_user
        assert meta.owner_groupname == expected_group
    else:
        # Fallback for Windows
        assert isinstance(meta.owner_username, str)
        assert isinstance(meta.owner_groupname, str)


def test_inspect_allocation():
    """
    Task 3: Allocation Metrics.
    Verifies apparent file size matches size_bytes, and allocated_size
    corresponds to st_blocks * 512 bytes.
    """
    test_file = TEST_ENV / "alloc_test.bin"
    with open(test_file, "wb") as f:
        f.write(b"A" * 1234)
        
    stat_info = test_file.stat()
    meta = get_metadata(test_file)
    
    assert meta.size == 1234
    
    # st_blocks is Unix-specific. On other platforms, fallback to size on disk
    if hasattr(stat_info, "st_blocks"):
        expected_blocks_size = stat_info.st_blocks * 512
        assert meta.allocated_size == expected_blocks_size
    else:
        assert meta.allocated_size >= 1234


def test_inspect_hard_links():
    """
    Task 4: Hard Link Count.
    Verifies that the hard link count changes when files are linked.
    """
    test_file = TEST_ENV / "link_test.bin"
    with open(test_file, "wb") as f:
        f.write(b"link")
        
    meta_single = get_metadata(test_file)
    assert meta_single.hard_links == 1
    
    # Create a hard link
    link_file = TEST_ENV / "link_test_hard.bin"
    if link_file.exists():
        link_file.unlink()
        
    try:
        os.link(str(test_file), str(link_file))
        meta_linked = get_metadata(test_file)
        assert meta_linked.hard_links == 2
    finally:
        if link_file.exists():
            link_file.unlink()


def test_inspect_inode_device():
    """
    Task 5: Inode & Device Identifiers.
    Verifies that the unique Inode number and Device ID match os.stat results.
    """
    test_file = TEST_ENV / "inode_test.bin"
    with open(test_file, "wb") as f:
        f.write(b"data")
        
    stat_info = test_file.stat()
    meta = get_metadata(test_file)
    
    assert meta.inode == stat_info.st_ino
    assert meta.device == stat_info.st_dev


def test_inspect_extended_attributes():
    """
    Task 6: Extended Attributes (xattr).
    Checks that user-defined extended attributes can be retrieved if supported
    by the platform and filesystem.
    """
    test_file = TEST_ENV / "xattr_test.bin"
    with open(test_file, "wb") as f:
        f.write(b"data")
        
    # Attempt to set an extended attribute (Linux/macOS specific)
    has_xattr_support = False
    if sys.platform in ("linux", "darwin"):
        try:
            import xattr
            # Use user namespace for custom attribute
            os.setxattr(str(test_file), "user.ctf_flag", b"FLAG{filesystem_metadata_ftw}")
            has_xattr_support = True
        except (ImportError, OSError, AttributeError):
            pass
            
    meta = get_metadata(test_file)
    
    if has_xattr_support:
        assert "user.ctf_flag" in meta.extended_attributes
        assert meta.extended_attributes["user.ctf_flag"] == "FLAG{filesystem_metadata_ftw}"
    else:
        # Should gracefully return an empty dict if not supported or none defined
        assert isinstance(meta.extended_attributes, dict)


# --- CLI Validation Tests ---

def test_forensics_inspect_cli_success():
    """
    Verifies that 'forensics inspect' successfully runs via the CLI and prints 
    the metadata in a formatted table.
    """
    runner = CliRunner()
    test_file = TEST_ENV / "test_file.png"
    with open(test_file, "wb") as f:
        f.write(b"\x89PNG\r\n\x1a\n")
        
    result = runner.invoke(forensics_group, ["inspect", str(test_file)])
    assert result.exit_code == 0
    assert "Metadata: test_file.png" in result.output
    assert "Size" in result.output
    assert "89504E47" in result.output
    assert "Detected Type" in result.output
    assert "PNG Image" in result.output

def test_root_inspect_cli_success():
    """
    Verifies that 'ctf inspect' successfully runs at the root level and prints 
    the metadata in a formatted table.
    """
    from ctf.main import cli
    runner = CliRunner()
    test_file = TEST_ENV / "test_file_root.png"
    with open(test_file, "wb") as f:
        f.write(b"\x89PNG\r\n\x1a\n")
        
    result = runner.invoke(cli, ["inspect", str(test_file)])
    assert result.exit_code == 0
    assert "Metadata: test_file_root.png" in result.output
    assert "PNG Image" in result.output

def test_nested_inspect_cli_success():
    """
    Verifies that 'ctf forensics inspect' successfully runs via the root CLI.
    """
    from ctf.main import cli
    runner = CliRunner()
    test_file = TEST_ENV / "test_file_nested.png"
    with open(test_file, "wb") as f:
        f.write(b"\x89PNG\r\n\x1a\n")
        
    result = runner.invoke(cli, ["forensics", "inspect", str(test_file)])
    assert result.exit_code == 0
    assert "Metadata: test_file_nested.png" in result.output
    assert "PNG Image" in result.output

def test_forensics_inspect_cli_missing_file():
    """
    Verifies that the CLI fails gracefully when a non-existent file path is specified.
    """
    runner = CliRunner()
    result = runner.invoke(forensics_group, ["inspect", "non_existent_file.png"])
    assert result.exit_code != 0
    assert "does not exist" in result.output

def test_forensics_signatures_cli():
    """
    Verifies that 'forensics signatures' runs successfully and displays
    the list of supported file signatures.
    """
    runner = CliRunner()
    result = runner.invoke(forensics_group, ["signatures"])
    assert result.exit_code == 0
    assert "Supported Magic Signatures" in result.output
    assert "PNG Image" in result.output
    assert "89504E47" in result.output

# (Old flag-detect tests removed)

def test_pure_exif_parsing():
    """
    Verifies that get_metadata successfully parses EXIF fields from a mock PNG.
    """
    import struct
    tiff_data = (
        b"II\x2a\x00\x08\x00\x00\x00"  # TIFF Header
        b"\x02\x00"                    # Num entries
        b"\x0f\x01\x02\x00\x05\x00\x00\x00\x26\x00\x00\x00"  # Entry 1 (Make: offset 38)
        b"\x10\x01\x02\x00\x06\x00\x00\x00\x2b\x00\x00\x00"  # Entry 2 (Model: offset 43)
        b"\x00\x00\x00\x00"            # Next IFD offset
        b"TEST\x00"                    # Make value
        b"MODEL\x00"                   # Model value
    )
    png_header = b"\x89PNG\r\n\x1a\n"
    exif_chunk_type = b"eXIf"
    exif_chunk_len = struct.pack(">I", len(tiff_data))
    exif_chunk_crc = b"\x00\x00\x00\x00"
    
    mock_png = png_header + exif_chunk_len + exif_chunk_type + tiff_data + exif_chunk_crc
    
    test_file = TEST_ENV / "mock_exif_image.png"
    with open(test_file, "wb") as f:
        f.write(mock_png)
        
    meta = get_metadata(test_file)
    assert meta.exif_data.get("Make") == "TEST"
    assert meta.exif_data.get("Model") == "MODEL"

def test_exif_cli_options():
    """
    Verifies CLI options -i/--inode and -e/--exif work as expected.
    """
    from ctf.main import cli
    runner = CliRunner()
    test_file = TEST_ENV / "mock_exif_image.png"
    
    # Verify mutual exclusion error
    result_err = runner.invoke(cli, ["inspect", "-i", "-e", str(test_file)])
    assert result_err.exit_code != 0
    assert "Cannot specify both" in result_err.output
    
    # Verify --inode hides EXIF data table
    result_inode = runner.invoke(cli, ["inspect", "-i", str(test_file)])
    assert result_inode.exit_code == 0
    assert "Metadata: mock_exif_image.png" in result_inode.output
    assert "EXIF Metadata" not in result_inode.output
    
    # Verify --exif hides POSIX table but displays EXIF data table
    result_exif = runner.invoke(cli, ["inspect", "-e", str(test_file)])
    assert result_exif.exit_code == 0
    assert "Metadata: mock_exif_image.png" not in result_exif.output
    assert "EXIF Metadata" in result_exif.output
    assert "mock_exif_image" in result_exif.output
    assert "Make" in result_exif.output
    assert "TEST" in result_exif.output

def test_jpeg_comment_extraction():
    """
    Verifies that get_metadata successfully parses a JPEG comment (COM marker).
    """
    import struct
    comment_text = b"picoCTF{test_comment_flag}"
    comment_len = len(comment_text) + 2
    mock_jpeg = (
        b"\xff\xd8"
        b"\xff\xfe"
        + struct.pack(">H", comment_len)
        + comment_text
        + b"\xff\xd9"
    )
    
    test_file = TEST_ENV / "mock_comment.jpg"
    with open(test_file, "wb") as f:
        f.write(mock_jpeg)
        
    meta = get_metadata(test_file)
    assert meta.comment == "picoCTF{test_comment_flag}"

def test_exif_cli_options_warning():
    """
    Verifies CLI warns when no EXIF or comment is found and -e option is used.
    """
    from ctf.main import cli
    runner = CliRunner()
    test_file = TEST_ENV / "clean_no_exif.png"
    with open(test_file, "wb") as f:
        f.write(b"\x89PNG\r\n\x1a\n")
        
    result = runner.invoke(cli, ["inspect", "-e", str(test_file)])
    assert result.exit_code == 0
    assert "No EXIF or comment metadata found" in result.output

def test_exif_cli_comment_only():
    """
    Verifies CLI displays comment when -e is used and only comment is found.
    """
    from ctf.main import cli
    import struct
    runner = CliRunner()
    comment_text = b"flag{comment_only}"
    comment_len = len(comment_text) + 2
    mock_jpeg = b"\xff\xd8\xff\xfe" + struct.pack(">H", comment_len) + comment_text
    
    test_file = TEST_ENV / "comment_only.jpg"
    with open(test_file, "wb") as f:
        f.write(mock_jpeg)
        
    result = runner.invoke(cli, ["inspect", "-e", str(test_file)])
    assert result.exit_code == 0
    assert "Comment" in result.output
    assert "flag{comment_only}" in result.output