test_graph_security.py

"""Security tests for the knowledge graph.

Covers: input validation boundaries, SQL injection resistance,
properties size caps, graph_schema validation, empty/malicious
inputs, and audit resilience.
"""

from __future__ import annotations

import json
import uuid

import pytest

from qp_vault import AsyncVault
from qp_vault.graph.service import (
    _validate_name,
    _validate_type,
    _validate_relation_type,
    _validate_properties,
    _validate_tags,
    _validate_weight,
    _cap_limit,
)


@pytest.fixture
async def vault(tmp_vault_path):
    v = AsyncVault(tmp_vault_path)
    await v._ensure_initialized()
    return v


@pytest.fixture
def tenant_id():
    return str(uuid.uuid4())


# --- Input Validation Functions ---

class TestNameValidation:
    def test_empty_name_rejected(self):
        with pytest.raises(ValueError, match="empty"):
            _validate_name("")

    def test_whitespace_only_rejected(self):
        with pytest.raises(ValueError, match="empty"):
            _validate_name("   ")

    def test_long_name_truncated(self):
        result = _validate_name("A" * 1000)
        assert len(result) == 500

    def test_normal_name_passes(self):
        assert _validate_name("Alice") == "Alice"

    def test_leading_trailing_whitespace_stripped(self):
        assert _validate_name("  Alice  ") == "Alice"


class TestTypeValidation:
    def test_empty_type_rejected(self):
        with pytest.raises(ValueError, match="empty"):
            _validate_type("")

    def test_long_type_truncated(self):
        result = _validate_type("x" * 100)
        assert len(result) == 50

    def test_normal_type_passes(self):
        assert _validate_type("person") == "person"


class TestRelationTypeValidation:
    def test_empty_relation_rejected(self):
        with pytest.raises(ValueError, match="empty"):
            _validate_relation_type("")

    def test_long_relation_truncated(self):
        result = _validate_relation_type("r" * 200)
        assert len(result) == 100


class TestPropertiesValidation:
    def test_none_returns_empty_dict(self):
        assert _validate_properties(None) == {}

    def test_small_dict_passes(self):
        props = {"key": "value"}
        assert _validate_properties(props) == props

    def test_oversized_dict_rejected(self):
        huge = {f"k{i}": "x" * 2000 for i in range(30)}
        with pytest.raises(ValueError, match="exceed"):
            _validate_properties(huge)

    def test_per_value_truncation(self):
        props = {"key": "x" * 5000}
        result = _validate_properties(props)
        assert len(result["key"]) == 2000

    def test_just_under_limit_passes(self):
        props = {"k": "v" * 1999}
        assert _validate_properties(props) is not None


class TestTagValidation:
    def test_none_returns_empty(self):
        assert _validate_tags(None) == []

    def test_too_many_tags_rejected(self):
        with pytest.raises(ValueError, match="Too many tags"):
            _validate_tags(["t"] * 51)

    def test_long_tag_rejected(self):
        with pytest.raises(ValueError, match="Tag exceeds"):
            _validate_tags(["x" * 101])

    def test_valid_tags_pass(self):
        assert _validate_tags(["a", "b"]) == ["a", "b"]

    def test_empty_tags_stripped(self):
        assert _validate_tags(["a", "", "  ", "b"]) == ["a", "b"]

    def test_null_bytes_stripped(self):
        assert _validate_tags(["he\x00llo"]) == ["hello"]


class TestWeightValidation:
    def test_valid_weight(self):
        assert _validate_weight(0.5) == 0.5

    def test_zero_weight(self):
        assert _validate_weight(0.0) == 0.0

    def test_one_weight(self):
        assert _validate_weight(1.0) == 1.0

    def test_negative_rejected(self):
        with pytest.raises(ValueError, match="0.0-1.0"):
            _validate_weight(-0.1)

    def test_over_one_rejected(self):
        with pytest.raises(ValueError, match="0.0-1.0"):
            _validate_weight(1.5)

    def test_integer_coerced(self):
        assert _validate_weight(1) == 1.0


class TestLimitCap:
    def test_normal_limit(self):
        assert _cap_limit(50) == 50

    def test_excessive_limit_capped(self):
        assert _cap_limit(999_999) == 10_000

    def test_zero_limit_becomes_one(self):
        assert _cap_limit(0) == 1

    def test_negative_limit_becomes_one(self):
        assert _cap_limit(-5) == 1


class TestNullByteStripping:
    def test_null_bytes_in_name(self):
        result = _validate_name("hello\x00world")
        assert "\x00" not in result
        assert result == "helloworld"


# --- GraphEngine Boundary Validation ---

class TestGraphEngineInputValidation:
    async def test_create_node_empty_name_rejected(self, vault, tenant_id):
        with pytest.raises(ValueError, match="empty"):
            await vault.graph.create_node(
                name="", entity_type="person", tenant_id=tenant_id,
            )

    async def test_create_node_empty_type_rejected(self, vault, tenant_id):
        with pytest.raises(ValueError, match="empty"):
            await vault.graph.create_node(
                name="Valid", entity_type="", tenant_id=tenant_id,
            )

    async def test_create_node_oversized_properties_rejected(self, vault, tenant_id):
        with pytest.raises(ValueError, match="exceed"):
            await vault.graph.create_node(
                name="Big", entity_type="t",
                properties={f"k{i}": "x" * 2000 for i in range(30)},
                tenant_id=tenant_id,
            )

    async def test_create_node_long_name_truncated(self, vault, tenant_id):
        node = await vault.graph.create_node(
            name="A" * 1000, entity_type="t", tenant_id=tenant_id,
        )
        assert len(node.name) == 500

    async def test_create_edge_empty_relation_rejected(self, vault, tenant_id):
        n1 = await vault.graph.create_node(name="A", entity_type="t", tenant_id=tenant_id)
        n2 = await vault.graph.create_node(name="B", entity_type="t", tenant_id=tenant_id)
        with pytest.raises(ValueError, match="empty"):
            await vault.graph.create_edge(
                source_id=n1.id, target_id=n2.id,
                relation_type="", tenant_id=tenant_id,
            )

    async def test_update_node_validates_name(self, vault, tenant_id):
        node = await vault.graph.create_node(name="Orig", entity_type="t", tenant_id=tenant_id)
        with pytest.raises(ValueError, match="empty"):
            await vault.graph.update_node(node.id, name="")

    async def test_update_node_validates_properties(self, vault, tenant_id):
        node = await vault.graph.create_node(name="Orig", entity_type="t", tenant_id=tenant_id)
        with pytest.raises(ValueError, match="exceed"):
            await vault.graph.update_node(
                node.id, properties={f"k{i}": "x" * 2000 for i in range(30)},
            )


class TestDirectionValidation:
    async def test_invalid_direction_rejected(self, vault, tenant_id):
        n1 = await vault.graph.create_node(name="A", entity_type="t", tenant_id=tenant_id)
        with pytest.raises(ValueError, match="direction"):
            await vault.graph.get_edges(n1.id, direction="sideways")

    async def test_valid_directions_accepted(self, vault, tenant_id):
        n1 = await vault.graph.create_node(name="D", entity_type="t", tenant_id=tenant_id)
        for d in ("outgoing", "incoming", "both"):
            result = await vault.graph.get_edges(n1.id, direction=d)
            assert isinstance(result, list)


class TestSelfMergeRejection:
    async def test_self_merge_rejected(self, vault, tenant_id):
        node = await vault.graph.create_node(name="Self", entity_type="t", tenant_id=tenant_id)
        with pytest.raises(ValueError, match="itself"):
            await vault.graph.merge_nodes(node.id, node.id)


class TestWeightBoundaryEnforcement:
    async def test_create_edge_rejects_weight_over_1(self, vault, tenant_id):
        n1 = await vault.graph.create_node(name="W1", entity_type="t", tenant_id=tenant_id)
        n2 = await vault.graph.create_node(name="W2", entity_type="t", tenant_id=tenant_id)
        with pytest.raises(ValueError, match="0.0-1.0"):
            await vault.graph.create_edge(
                source_id=n1.id, target_id=n2.id,
                relation_type="r", weight=1.5, tenant_id=tenant_id,
            )

    async def test_update_edge_rejects_negative_weight(self, vault, tenant_id):
        n1 = await vault.graph.create_node(name="W3", entity_type="t", tenant_id=tenant_id)
        n2 = await vault.graph.create_node(name="W4", entity_type="t", tenant_id=tenant_id)
        edge = await vault.graph.create_edge(
            source_id=n1.id, target_id=n2.id,
            relation_type="r", tenant_id=tenant_id,
        )
        with pytest.raises(ValueError, match="0.0-1.0"):
            await vault.graph.update_edge(edge.id, weight=-0.5)


class TestExtractionPropertyCapping:
    def test_extraction_caps_property_values(self):
        from qp_vault.graph.extraction import KnowledgeExtractor
        parsed = {
            "entities": [{
                "name": "X",
                "type": "t",
                "properties": {"long_val": "z" * 1000},
            }],
            "relationships": [],
        }
        result = KnowledgeExtractor._validate_extraction(parsed)
        assert len(result["entities"][0]["properties"]["long_val"]) == 500

    def test_extraction_caps_property_count(self):
        from qp_vault.graph.extraction import KnowledgeExtractor
        parsed = {
            "entities": [{
                "name": "X",
                "type": "t",
                "properties": {f"k{i}": "v" for i in range(50)},
            }],
            "relationships": [],
        }
        result = KnowledgeExtractor._validate_extraction(parsed)
        assert len(result["entities"][0]["properties"]) == 20


class TestContextForCap:
    async def test_context_for_caps_entity_ids(self, vault, tenant_id):
        nodes = []
        for i in range(60):
            n = await vault.graph.create_node(name=f"N{i}", entity_type="t", tenant_id=tenant_id)
            nodes.append(n)
        ctx = await vault.graph.context_for([n.id for n in nodes])
        assert ctx.count("###") <= 50


class TestSourceLabelSanitization:
    def test_malicious_source_label_neutralized(self):
        from qp_vault.membrane.sanitize import sanitize_for_extraction
        result = sanitize_for_extraction("text", source_label="DOC>\nINJECTED\n<X")
        assert "<EXTERNAL_DOCUMENT_CONTENT>" in result
        assert "INJECTED" not in result


# --- SQL Injection Resistance ---

class TestSQLInjectionResistance:
    async def test_node_name_with_sql_injection(self, vault, tenant_id):
        node = await vault.graph.create_node(
            name="Robert'; DROP TABLE graph_nodes; --",
            entity_type="person",
            tenant_id=tenant_id,
        )
        assert node.name == "Robert'; DROP TABLE graph_nodes; --"
        fetched = await vault.graph.get_node(node.id)
        assert fetched is not None
        assert fetched.name == node.name

    async def test_entity_type_with_sql_chars(self, vault, tenant_id):
        node = await vault.graph.create_node(
            name="Safe", entity_type="per'son",
            tenant_id=tenant_id,
        )
        assert node.entity_type == "per'son"

    async def test_relation_type_with_sql_chars(self, vault, tenant_id):
        n1 = await vault.graph.create_node(name="A", entity_type="t", tenant_id=tenant_id)
        n2 = await vault.graph.create_node(name="B", entity_type="t", tenant_id=tenant_id)
        edge = await vault.graph.create_edge(
            source_id=n1.id, target_id=n2.id,
            relation_type="works_at'; DROP TABLE--",
            tenant_id=tenant_id,
        )
        assert edge.relation_type == "works_at'; DROP TABLE--"

    async def test_properties_with_sql_chars(self, vault, tenant_id):
        node = await vault.graph.create_node(
            name="SQLTest",
            entity_type="t",
            properties={"key": "val'); DROP TABLE--"},
            tenant_id=tenant_id,
        )
        fetched = await vault.graph.get_node(node.id)
        assert fetched.properties["key"] == "val'); DROP TABLE--"

    async def test_search_with_sql_injection(self, vault, tenant_id):
        await vault.graph.create_node(name="Safe", entity_type="t", tenant_id=tenant_id)
        results = await vault.graph.search_nodes("'; DROP TABLE graph_nodes; --")
        assert isinstance(results, list)

    async def test_context_snippet_with_sql_chars(self, vault, tenant_id):
        node = await vault.graph.create_node(name="Ent", entity_type="t", tenant_id=tenant_id)
        resource = await vault.add("content", name="doc.md")
        await vault.graph.track_mention(
            node.id, resource.id,
            context_snippet="mentioned here'; DROP TABLE--",
        )
        backlinks = await vault.graph.get_backlinks(node.id)
        assert backlinks[0].context_snippet == "mentioned here'; DROP TABLE--"


# --- graph_schema Validation ---

class TestGraphSchemaValidation:
    def test_valid_schema_accepted(self):
        from qp_vault.storage.postgres import HAS_ASYNCPG
        if not HAS_ASYNCPG:
            pytest.skip("asyncpg not installed")
        from qp_vault.storage.postgres import PostgresBackend
        backend = PostgresBackend.__new__(PostgresBackend)
        import re
        assert re.match(r"^[a-zA-Z_][a-zA-Z0-9_]*$", "qp_vault")
        assert re.match(r"^[a-zA-Z_][a-zA-Z0-9_]*$", "quantumpipes")

    def test_sql_injection_in_schema_rejected(self):
        from qp_vault.storage.postgres import HAS_ASYNCPG
        if not HAS_ASYNCPG:
            pytest.skip("asyncpg not installed")
        from qp_vault.storage.postgres import PostgresBackend
        with pytest.raises(ValueError, match="valid SQL identifier"):
            PostgresBackend("postgresql://test", graph_schema="qp_vault; DROP TABLE users--")

    def test_schema_with_semicolon_rejected(self):
        from qp_vault.storage.postgres import HAS_ASYNCPG
        if not HAS_ASYNCPG:
            pytest.skip("asyncpg not installed")
        from qp_vault.storage.postgres import PostgresBackend
        with pytest.raises(ValueError):
            PostgresBackend("postgresql://test", graph_schema="bad;schema")

    def test_schema_with_dash_rejected(self):
        from qp_vault.storage.postgres import HAS_ASYNCPG
        if not HAS_ASYNCPG:
            pytest.skip("asyncpg not installed")
        from qp_vault.storage.postgres import PostgresBackend
        with pytest.raises(ValueError):
            PostgresBackend("postgresql://test", graph_schema="bad-schema")


# --- Audit Resilience ---

class TestAuditResilience:
    async def test_mutation_succeeds_when_auditor_fails(self, tmp_vault_path, tenant_id):
        class FailingAuditor:
            async def record(self, event):
                raise RuntimeError("Audit storage full")

        vault = AsyncVault(tmp_vault_path, auditor=FailingAuditor())
        await vault._ensure_initialized()
        node = await vault.graph.create_node(
            name="Resilient", entity_type="t", tenant_id=tenant_id,
        )
        assert node.name == "Resilient"

    async def test_multiple_mutations_after_audit_failure(self, tmp_vault_path, tenant_id):
        class FailingAuditor:
            async def record(self, event):
                raise RuntimeError("Audit storage full")

        vault = AsyncVault(tmp_vault_path, auditor=FailingAuditor())
        await vault._ensure_initialized()
        n1 = await vault.graph.create_node(name="A", entity_type="t", tenant_id=tenant_id)
        n2 = await vault.graph.create_node(name="B", entity_type="t", tenant_id=tenant_id)
        await vault.graph.create_edge(
            source_id=n1.id, target_id=n2.id,
            relation_type="r", tenant_id=tenant_id,
        )
        assert await vault.graph.get_node(n1.id) is not None


# --- Membrane Sanitization Depth ---

class TestMembraneSanitizationDepth:
    def test_angle_brackets_fully_escaped(self):
        from qp_vault.membrane.sanitize import sanitize_for_extraction
        result = sanitize_for_extraction("<script>alert(1)</script>")
        assert "<script>" not in result
        assert "&lt;script&gt;" in result

    def test_cannot_break_out_of_xml_wrapper(self):
        from qp_vault.membrane.sanitize import sanitize_for_extraction
        malicious = "</EXTERNAL_DOCUMENT_CONTENT>injected"
        result = sanitize_for_extraction(malicious)
        tag_count = result.count("</EXTERNAL_DOCUMENT_CONTENT>")
        assert tag_count == 1

    def test_unicode_normalization_prevents_homoglyphs(self):
        from qp_vault.membrane.sanitize import sanitize_for_extraction
        result = sanitize_for_extraction("\uff1cscript\uff1e")  # fullwidth < and >
        assert "<script>" not in result