120 lines
3.9 KiB
Python
120 lines
3.9 KiB
Python
"""Adversarial tests for the pre-parser content security gate."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import io
|
|
import zipfile
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
|
|
from ragcore.application.security import (
|
|
ContentSecurityRequest,
|
|
QuarantinedContentError,
|
|
QuarantineReason,
|
|
SecurityDisposition,
|
|
)
|
|
from ragcore.domain.audit import AuditEventType
|
|
from ragcore.infrastructure.security import ContentSecurityGateService, ContentSecurityPolicy
|
|
|
|
|
|
class RecordingAudit:
|
|
def __init__(self) -> None:
|
|
self.events: list[dict[str, object]] = []
|
|
|
|
async def record(self, **event: object) -> None:
|
|
self.events.append(event)
|
|
|
|
|
|
def request(
|
|
content: bytes, *, media_type: str = "text/plain", name: str = "note.txt"
|
|
) -> ContentSecurityRequest:
|
|
return ContentSecurityRequest(
|
|
workspace_id=uuid4(),
|
|
source_id=uuid4(),
|
|
content=content,
|
|
declared_media_type=media_type,
|
|
original_name=name,
|
|
trace_id="trace-content-security",
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mime_mismatch_and_disallowed_type_are_quarantined() -> None:
|
|
audit = RecordingAudit()
|
|
gate = ContentSecurityGateService(audit)
|
|
|
|
mismatch = await gate.inspect(request(b"%PDF-1.7", media_type="text/plain", name="note.txt"))
|
|
disallowed = await gate.inspect(
|
|
request(b"#!/bin/sh\necho unsafe", media_type="text/x-shellscript", name="run.sh")
|
|
)
|
|
|
|
assert mismatch.disposition is SecurityDisposition.QUARANTINED
|
|
assert mismatch.reason is QuarantineReason.MIME_MISMATCH
|
|
assert disallowed.reason is QuarantineReason.DISALLOWED_MEDIA_TYPE
|
|
with pytest.raises(QuarantinedContentError):
|
|
mismatch.release_for_parser()
|
|
assert len(audit.events) == 2
|
|
assert all(event["event_type"] is AuditEventType.SECURITY_FAILURE for event in audit.events)
|
|
|
|
|
|
def _zip_with_entry(name: str, content: bytes, *, compression: int = zipfile.ZIP_DEFLATED) -> bytes:
|
|
output = io.BytesIO()
|
|
with zipfile.ZipFile(output, "w", compression=compression) as archive:
|
|
archive.writestr(name, content)
|
|
return output.getvalue()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_archive_ratio_and_size_limits_quarantine_before_release() -> None:
|
|
audit = RecordingAudit()
|
|
ratio_gate = ContentSecurityGateService(
|
|
audit,
|
|
policy=ContentSecurityPolicy(max_archive_ratio=2.0, max_archive_uncompressed_bytes=256),
|
|
)
|
|
size_gate = ContentSecurityGateService(
|
|
audit,
|
|
policy=ContentSecurityPolicy(max_archive_ratio=2.0, max_archive_uncompressed_bytes=32),
|
|
)
|
|
|
|
ratio = await ratio_gate.inspect(
|
|
request(
|
|
_zip_with_entry("data.txt", b"a" * 128), media_type="application/zip", name="data.zip"
|
|
)
|
|
)
|
|
size = await size_gate.inspect(
|
|
request(
|
|
_zip_with_entry("data.txt", b"b" * 40, compression=zipfile.ZIP_STORED),
|
|
media_type="application/zip",
|
|
name="data.zip",
|
|
)
|
|
)
|
|
|
|
assert ratio.reason is QuarantineReason.ARCHIVE_RATIO_LIMIT
|
|
assert size.reason is QuarantineReason.ARCHIVE_SIZE_LIMIT
|
|
for decision in (ratio, size):
|
|
with pytest.raises(QuarantinedContentError):
|
|
decision.release_for_parser()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_archive_depth_limit_and_quarantine_cannot_reach_parser_or_index() -> None:
|
|
audit = RecordingAudit()
|
|
policy = ContentSecurityPolicy(max_archive_depth=1, max_archive_ratio=10_000.0)
|
|
gate = ContentSecurityGateService(audit, policy=policy)
|
|
nested = _zip_with_entry("inner.zip", _zip_with_entry("note.txt", b"safe"))
|
|
|
|
decision = await gate.inspect(request(nested, media_type="application/zip", name="outer.zip"))
|
|
|
|
assert decision.reason is QuarantineReason.ARCHIVE_DEPTH_LIMIT
|
|
parser_calls = 0
|
|
|
|
def parser_or_index(content: bytes) -> None:
|
|
nonlocal parser_calls
|
|
parser_calls += 1
|
|
assert content
|
|
|
|
with pytest.raises(QuarantinedContentError):
|
|
parser_or_index(decision.release_for_parser())
|
|
assert parser_calls == 0
|