Agent/tests/test_dream_security.py

import asyncio
import json
import os
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch

from cyber_agent.core.dream import (
    DreamScope,
    DreamScopeError,
    cross_trace_integrate,
)
from cyber_agent.core.memory import (
    MEMORY_IDENTITY_CONTEXT_KEY,
    MemoryConfig,
    MemoryPathError,
    compute_memory_identity,
    is_declared_memory_path,
    load_memory_files,
    resolve_memory_path,
)
from cyber_agent.tools.builtin.memory import dream as dream_tool
from cyber_agent.trace.models import Trace


class FakeDreamStore:
    def __init__(self, base_path: Path, traces):
        self.base_path = base_path
        self.traces = {trace.trace_id: trace for trace in traces}
        self.base_path.mkdir(parents=True, exist_ok=True)

    async def list_traces(self, limit=1000):
        return list(self.traces.values())[:limit]

    async def get_trace(self, trace_id):
        return self.traces.get(trace_id)

    async def get_cognition_log(self, trace_id):
        path = self._get_cognition_log_file(trace_id)
        if not path.exists():
            return {"trace_id": trace_id, "events": []}
        return json.loads(path.read_text(encoding="utf-8"))

    async def mark_reflections_consumed(
        self, trace_id, reflections, consumed_at
    ):
        log = await self.get_cognition_log(trace_id)
        wanted = {
            (tuple(item.get("sequence_range") or []), item.get("summary", ""))
            for item in reflections
        }
        changed = 0
        for event in log.get("events", []):
            key = (
                tuple(event.get("sequence_range") or []),
                event.get("summary", ""),
            )
            if (
                event.get("type") == "reflection"
                and not event.get("consumed_at")
                and key in wanted
            ):
                event["consumed_at"] = consumed_at
                changed += 1
        self._get_cognition_log_file(trace_id).write_text(
            json.dumps(log),
            encoding="utf-8",
        )
        return changed

    def _get_cognition_log_file(self, trace_id):
        path = self.base_path / trace_id / "cognition_log.json"
        path.parent.mkdir(parents=True, exist_ok=True)
        return path

    def set_log(self, trace_id, events):
        self._get_cognition_log_file(trace_id).write_text(
            json.dumps({"trace_id": trace_id, "events": events}),
            encoding="utf-8",
        )


def make_trace(trace_id, scope, *, uid=None, agent_type=None, identity=None):
    return Trace(
        trace_id=trace_id,
        mode="agent",
        uid=scope.uid if uid is None else uid,
        agent_type=scope.agent_type if agent_type is None else agent_type,
        context={
            MEMORY_IDENTITY_CONTEXT_KEY: (
                scope.memory_identity if identity is None else identity
            )
        },
    )


def reflection(summary):
    return {
        "type": "reflection",
        "timestamp": f"ts-{summary}",
        "sequence_range": [1, 2],
        "summary": summary,
    }


class MemorySecurityTests(unittest.TestCase):
    def test_identity_uses_real_base_and_file_patterns(self):
        with tempfile.TemporaryDirectory() as tmp:
            base = Path(tmp) / "memory"
            base.mkdir()
            alias = Path(tmp) / "memory-alias"
            alias.symlink_to(base, target_is_directory=True)

            first = MemoryConfig(
                base_path=str(base),
                files={"journals/**/*.md": "journals", "identity.md": "self"},
                dream_prompt="first prompt",
            )
            second = MemoryConfig(
                base_path=str(alias),
                files={"identity.md": "different purpose", "journals/**/*.md": ""},
                dream_prompt="different prompt",
            )
            self.assertEqual(
                compute_memory_identity(first), compute_memory_identity(second)
            )

            third = MemoryConfig(
                base_path=str(base),
                files={"identity.md": "self"},
            )
            self.assertNotEqual(
                compute_memory_identity(first), compute_memory_identity(third)
            )

    def test_declared_glob_does_not_let_star_cross_directories(self):
        config = MemoryConfig(
            base_path="/tmp/not-used",
            files={"relationships/*.md": "people", "journals/**/*.md": "logs"},
        )
        self.assertTrue(is_declared_memory_path(config, "relationships/a.md"))
        self.assertFalse(
            is_declared_memory_path(config, "relationships/private/a.md")
        )
        self.assertTrue(is_declared_memory_path(config, "journals/2026/07/a.md"))

    def test_resolution_rejects_non_normal_paths_and_symlink_escape(self):
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            base = root / "memory"
            outside = root / "outside"
            base.mkdir()
            outside.mkdir()
            (base / "link").symlink_to(outside, target_is_directory=True)
            config = MemoryConfig(base_path=str(base), files={"**/*.md": "all"})

            for unsafe in (
                "../outside/escape.md",
                "./identity.md",
                "/tmp/escape.md",
                "nested//identity.md",
                "nested\\identity.md",
                "link/escape.md",
            ):
                with self.subTest(path=unsafe):
                    with self.assertRaises(MemoryPathError):
                        resolve_memory_path(config, unsafe)

    def test_loader_does_not_follow_symlink_outside_base(self):
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            base = root / "memory"
            outside = root / "outside"
            base.mkdir()
            outside.mkdir()
            (outside / "secret.md").write_text("secret", encoding="utf-8")
            (base / "safe.md").write_text("safe", encoding="utf-8")
            (base / "linked.md").symlink_to(outside / "secret.md")
            config = MemoryConfig(base_path=str(base), files={"*.md": "root"})

            loaded = load_memory_files(config)
            self.assertEqual(loaded, [("safe.md", "root", "safe")])


class DreamIsolationTests(unittest.IsolatedAsyncioTestCase):
    async def asyncSetUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.root = Path(self.tmp.name)
        self.memory_root = self.root / "memory"
        self.memory_root.mkdir()
        self.config = MemoryConfig(
            base_path=str(self.memory_root),
            files={"identity.md": "self", "journals/**/*.md": "logs"},
        )
        self.scope = DreamScope(
            uid="user-1",
            agent_type="writer",
            memory_identity=compute_memory_identity(self.config),
        )

    async def asyncTearDown(self):
        self.tmp.cleanup()

    async def test_filters_uid_agent_type_and_memory_identity_together(self):
        matching = make_trace("matching", self.scope)
        wrong_uid = make_trace("wrong-uid", self.scope, uid="user-2")
        wrong_agent = make_trace("wrong-agent", self.scope, agent_type="researcher")
        wrong_memory = make_trace("wrong-memory", self.scope, identity="memory-v1:other")
        store = FakeDreamStore(
            self.root / "traces",
            [matching, wrong_uid, wrong_agent, wrong_memory],
        )
        for trace in (matching, wrong_uid, wrong_agent, wrong_memory):
            store.set_log(trace.trace_id, [reflection(trace.trace_id)])

        calls = []

        async def llm_call(**kwargs):
            calls.append(kwargs)
            return {"content": '{"updates": [], "reasoning": "merged"}'}

        consumed, updated, reasoning = await cross_trace_integrate(
            store, llm_call, self.config, self.scope
        )

        self.assertEqual((consumed, updated, reasoning), (1, [], "merged"))
        prompt = calls[0]["messages"][1]["content"]
        self.assertIn("matching", prompt)
        self.assertNotIn("wrong-uid", prompt)
        self.assertNotIn("wrong-agent", prompt)
        self.assertNotIn("wrong-memory", prompt)
        self.assertTrue(
            (await store.get_cognition_log("matching"))["events"][0].get(
                "consumed_at"
            )
        )
        for trace_id in ("wrong-uid", "wrong-agent", "wrong-memory"):
            self.assertNotIn(
                "consumed_at",
                (await store.get_cognition_log(trace_id))["events"][0],
            )

    async def test_concurrent_dreams_keep_memory_configs_isolated(self):
        other_root = self.root / "other-memory"
        other_root.mkdir()
        other_config = MemoryConfig(
            base_path=str(other_root),
            files={"identity.md": "self"},
        )
        other_scope = DreamScope(
            uid=self.scope.uid,
            agent_type=self.scope.agent_type,
            memory_identity=compute_memory_identity(other_config),
        )
        first_trace = make_trace("first", self.scope)
        second_trace = make_trace("second", other_scope)
        store = FakeDreamStore(
            self.root / "traces", [first_trace, second_trace]
        )
        store.set_log(first_trace.trace_id, [reflection("first-summary")])
        store.set_log(second_trace.trace_id, [reflection("second-summary")])

        async def first_llm(**kwargs):
            await asyncio.sleep(0)
            self.assertIn("first-summary", kwargs["messages"][1]["content"])
            self.assertNotIn("second-summary", kwargs["messages"][1]["content"])
            return {
                "content": json.dumps(
                    {
                        "updates": [
                            {"path": "identity.md", "new_content": "first"}
                        ],
                        "reasoning": "first",
                    }
                )
            }

        async def second_llm(**kwargs):
            await asyncio.sleep(0)
            self.assertIn("second-summary", kwargs["messages"][1]["content"])
            self.assertNotIn("first-summary", kwargs["messages"][1]["content"])
            return {
                "content": json.dumps(
                    {
                        "updates": [
                            {"path": "identity.md", "new_content": "second"}
                        ],
                        "reasoning": "second",
                    }
                )
            }

        first_result, second_result = await asyncio.gather(
            cross_trace_integrate(
                store, first_llm, self.config, self.scope
            ),
            cross_trace_integrate(
                store, second_llm, other_config, other_scope
            ),
        )

        self.assertEqual(first_result[0], 1)
        self.assertEqual(second_result[0], 1)
        self.assertEqual(
            (self.memory_root / "identity.md").read_text(encoding="utf-8"),
            "first",
        )
        self.assertEqual(
            (other_root / "identity.md").read_text(encoding="utf-8"),
            "second",
        )

    async def test_scope_identity_must_match_memory_config(self):
        trace = make_trace("trace", self.scope)
        store = FakeDreamStore(self.root / "traces", [trace])
        store.set_log(trace.trace_id, [reflection("summary")])
        wrong_scope = DreamScope(
            uid=self.scope.uid,
            agent_type=self.scope.agent_type,
            memory_identity="memory-v1:wrong",
        )

        async def llm_call(**kwargs):
            self.fail("LLM must not run for a mismatched scope")

        with self.assertRaises(DreamScopeError):
            await cross_trace_integrate(
                store, llm_call, self.config, wrong_scope
            )

    async def test_allowed_updates_are_written_then_reflection_is_consumed(self):
        trace = make_trace("trace", self.scope)
        store = FakeDreamStore(self.root / "traces", [trace])
        store.set_log(trace.trace_id, [reflection("summary")])

        async def llm_call(**kwargs):
            return {
                "content": json.dumps(
                    {
                        "updates": [
                            {"path": "identity.md", "new_content": "new identity"},
                            {
                                "path": "journals/2026/entry.md",
                                "new_content": "new journal",
                            },
                        ],
                        "reasoning": "useful",
                    }
                )
            }

        result = await cross_trace_integrate(
            store, llm_call, self.config, self.scope
        )

        self.assertEqual(
            result, (1, ["identity.md", "journals/2026/entry.md"], "useful")
        )
        self.assertEqual(
            (self.memory_root / "identity.md").read_text(encoding="utf-8"),
            "new identity",
        )
        self.assertEqual(
            (self.memory_root / "journals/2026/entry.md").read_text(
                encoding="utf-8"
            ),
            "new journal",
        )
        self.assertTrue(
            (await store.get_cognition_log(trace.trace_id))["events"][0].get(
                "consumed_at"
            )
        )
        self.assertEqual(list(self.memory_root.rglob("*.tmp")), [])

    async def test_invalid_update_rejects_entire_plan_without_consuming(self):
        trace = make_trace("trace", self.scope)
        store = FakeDreamStore(self.root / "traces", [trace])
        unsafe_paths = [
            "../memory-evil/escape.md",
            "/tmp/escape.md",
            "not-declared.txt",
            "journals/../escape.md",
            "journals\\escape.md",
        ]

        for index, unsafe_path in enumerate(unsafe_paths):
            with self.subTest(path=unsafe_path):
                event = reflection(f"summary-{index}")
                store.set_log(trace.trace_id, [event])

                async def llm_call(**kwargs):
                    return {
                        "content": json.dumps(
                            {
                                "updates": [
                                    {
                                        "path": unsafe_path,
                                        "new_content": "must not write",
                                    }
                                ],
                                "reasoning": "bad",
                            }
                        )
                    }

                result = await cross_trace_integrate(
                    store, llm_call, self.config, self.scope
                )
                self.assertEqual(result, (0, [], ""))
                self.assertNotIn(
                    "consumed_at",
                    (await store.get_cognition_log(trace.trace_id))["events"][0],
                )

        self.assertFalse((self.root / "memory-evil" / "escape.md").exists())

    async def test_symlink_escape_is_rejected_without_consuming(self):
        outside = self.root / "outside"
        outside.mkdir()
        (self.memory_root / "journals").symlink_to(
            outside, target_is_directory=True
        )
        trace = make_trace("trace", self.scope)
        store = FakeDreamStore(self.root / "traces", [trace])
        store.set_log(trace.trace_id, [reflection("summary")])

        async def llm_call(**kwargs):
            return {
                "content": json.dumps(
                    {
                        "updates": [
                            {
                                "path": "journals/escape.md",
                                "new_content": "must not escape",
                            }
                        ],
                        "reasoning": "bad",
                    }
                )
            }

        result = await cross_trace_integrate(
            store, llm_call, self.config, self.scope
        )
        self.assertEqual(result, (0, [], ""))
        self.assertFalse((outside / "escape.md").exists())
        self.assertNotIn(
            "consumed_at",
            (await store.get_cognition_log(trace.trace_id))["events"][0],
        )

    async def test_failed_atomic_replace_preserves_old_file_and_does_not_consume(self):
        target = self.memory_root / "identity.md"
        target.write_text("old identity", encoding="utf-8")
        trace = make_trace("trace", self.scope)
        store = FakeDreamStore(self.root / "traces", [trace])
        store.set_log(trace.trace_id, [reflection("summary")])

        async def llm_call(**kwargs):
            return {
                "content": json.dumps(
                    {
                        "updates": [
                            {"path": "identity.md", "new_content": "new identity"}
                        ],
                        "reasoning": "useful",
                    }
                )
            }

        with patch("cyber_agent.core.dream.os.replace", side_effect=OSError("boom")):
            with self.assertRaises(OSError):
                await cross_trace_integrate(
                    store, llm_call, self.config, self.scope
                )

        self.assertEqual(target.read_text(encoding="utf-8"), "old identity")
        self.assertNotIn(
            "consumed_at",
            (await store.get_cognition_log(trace.trace_id))["events"][0],
        )
        self.assertEqual(list(self.memory_root.glob(".identity.md.*.tmp")), [])

    async def test_multi_file_failure_rolls_back_every_published_update(self):
        identity = self.memory_root / "identity.md"
        journal = self.memory_root / "journals" / "today.md"
        journal.parent.mkdir()
        identity.write_text("old identity", encoding="utf-8")
        journal.write_text("old journal", encoding="utf-8")
        trace = make_trace("trace", self.scope)
        store = FakeDreamStore(self.root / "traces", [trace])
        store.set_log(trace.trace_id, [reflection("summary")])

        async def llm_call(**kwargs):
            return {
                "content": json.dumps({
                    "updates": [
                        {"path": "identity.md", "new_content": "new identity"},
                        {"path": "journals/today.md", "new_content": "new journal"},
                    ],
                    "reasoning": "useful",
                })
            }

        real_replace = os.replace
        replace_calls = 0

        def fail_second_target(source, destination):
            nonlocal replace_calls
            replace_calls += 1
            if replace_calls == 2:
                raise OSError("second replacement failed")
            return real_replace(source, destination)

        with patch(
            "cyber_agent.core.dream.os.replace",
            side_effect=fail_second_target,
        ):
            with self.assertRaisesRegex(OSError, "second replacement failed"):
                await cross_trace_integrate(
                    store, llm_call, self.config, self.scope
                )

        self.assertEqual("old identity", identity.read_text(encoding="utf-8"))
        self.assertEqual("old journal", journal.read_text(encoding="utf-8"))
        self.assertNotIn(
            "consumed_at",
            (await store.get_cognition_log(trace.trace_id))["events"][0],
        )
        self.assertEqual(
            [],
            list(self.memory_root.rglob("*.dream-stage")),
        )

    async def test_builtin_dream_requires_memory_config_and_scope_in_context(self):
        class Runner:
            trace_store = object()
            llm_call = object()

        missing = await dream_tool(context={"runner": Runner()})
        self.assertEqual(missing.error, "memory_config not in tool context")

        missing_scope = await dream_tool(
            context={"runner": Runner(), "memory_config": self.config}
        )
        self.assertEqual(missing_scope.error, "dream_scope not in tool context")


if __name__ == "__main__":
    unittest.main()