# ========= Copyright 2026 @ Strukto.AI All Rights Reserved. =========
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ========= Copyright 2026 @ Strukto.AI All Rights Reserved. =========

import errno

import pytest

from mirage import Action, CommandContext, Deny, Policy, Workspace
from mirage.commands.builtin.utils.limit import LimitExceededError
from mirage.io import IOResult
from mirage.policy import (CommandRule, ExecuteResultContext, OpsContext,
                           OpsResultContext, PolicyError)
from mirage.resource.ram import RAMResource
from mirage.types import Limit, MountMode, OnExceed
from mirage.workspace.session import SessionProfile

from mirage.workspace.session.permissions import (  # isort: skip
    CommandsBlock, PathsBlock)


def _role(**blocks) -> dict[str, SessionProfile]:
    """The workspace's default role, spelled as one document.

    Permissions live in exactly one place now, so what these tests used
    to pass as `permissions=` is `profiles.default`, which shapes the
    workspace's own session and every session created without a name.
    """
    return {"default": SessionProfile(**blocks)}


class NoInterpreters(Policy):

    async def pre_command(self, ctx: CommandContext) -> Action | None:
        if ctx.command == "python3":
            return Deny("interpreters are off")
        return None


@pytest.mark.asyncio
async def test_workspace_guards_refuse_before_backend_io():
    ws = Workspace(
        {"/data/": RAMResource()},
        mode=MountMode.WRITE,
        profiles=_role(commands=CommandsBlock(
            deny=(CommandRule(reason="production data is protected",
                              commands=("rm", ),
                              paths=("/data/prod/*", )), ))),
    )
    try:
        await ws.execute("mkdir -p /data/prod")
        await ws.ops.write("/data/prod/x.txt", b"keep\n")
        result = await ws.execute("rm /data/prod/x.txt")
        assert result.exit_code == 1
        assert result.stderr == (b"rm: /data/prod/x.txt: "
                                 b"production data is protected\n")
        out = await ws.execute("cat /data/prod/x.txt")
        assert out.stdout == b"keep\n"
        ok = await ws.execute("rm -f /data/prod/../other.txt 2>/dev/null; "
                              "echo done")
        assert b"done" in ok.stdout
    finally:
        await ws.close()


@pytest.mark.asyncio
async def test_policies_add_wins_over_runtime_placement():
    # python3 is runtime-bound in the default world; the pre_command
    # hook fires ahead of runtime resolution, so the refusal wins.
    ws = Workspace({"/data/": RAMResource()}, mode=MountMode.WRITE)
    try:
        ws.policies.add(NoInterpreters())
        result = await ws.execute("python3 -c 'print(1)'")
        # A whole-command refusal is bash's "found but may not run".
        assert result.exit_code == 126
        assert result.stderr == (
            b"python3: policy denied: interpreters are off\n")
    finally:
        await ws.close()


@pytest.mark.asyncio
async def test_policies_constructor_param_accepts_instances():
    ws = Workspace({"/data/": RAMResource()},
                   mode=MountMode.WRITE,
                   policies=[NoInterpreters()])
    try:
        result = await ws.execute("python3 -c 'print(1)'")
        assert result.exit_code == 126
        assert result.stderr == (
            b"python3: policy denied: interpreters are off\n")
    finally:
        await ws.close()


@pytest.mark.asyncio
async def test_guards_cover_shell_builtins_and_namespace_routes():
    # source is a dispatch-level shell builtin and touch is
    # namespace-routed; neither reaches handle_command, so this pins
    # the hook at the dispatch chokepoint.
    ws = Workspace(
        {"/data/": RAMResource()},
        mode=MountMode.WRITE,
        profiles=_role(commands=CommandsBlock(deny=(
            CommandRule(reason="disabled", commands=("source", )),
            CommandRule(reason="frozen",
                        commands=("touch", ),
                        paths=("/data/prod/*", )),
        ))),
    )
    try:
        result = await ws.execute("source /data/setup.sh")
        assert result.exit_code == 126
        assert result.stderr == b"source: policy denied: disabled\n"
        result = await ws.execute("touch /data/prod/x")
        assert result.exit_code == 1
        assert b"frozen" in result.stderr
        ok = await ws.execute("touch /data/dev-x && echo done")
        assert b"done" in ok.stdout
    finally:
        await ws.close()


@pytest.mark.asyncio
async def test_guards_cover_path_valued_flags():
    # shuf discovers its output path from -o, not a positional operand;
    # the policy context must include flag-valued paths.
    ws = Workspace(
        {"/data/": RAMResource()},
        mode=MountMode.WRITE,
        profiles=_role(commands=CommandsBlock(
            deny=(CommandRule(reason="prod is protected",
                              commands=("shuf", ),
                              paths=("/data/prod/*", )), ))),
    )
    try:
        await ws.execute("mkdir -p /data/prod")
        result = await ws.execute("shuf -e a -o /data/prod/out")
        assert result.exit_code == 1
        assert b"prod is protected" in result.stderr
        listing = await ws.execute("ls /data/prod")
        assert b"out" not in listing.stdout
    finally:
        await ws.close()


class ReadOnlyProd(Policy):

    async def pre_ops(self, ctx: OpsContext) -> Action | None:
        if ctx.write and ctx.path.virtual.startswith("/data/prod/"):
            return Deny("prod is read-only")
        return None


@pytest.mark.asyncio
async def test_path_guards_hold_at_the_programmatic_door():
    # ws.ops is the same seam FUSE comes through; a path-only guard
    # must refuse it, not just shell commands (#675).
    ws = Workspace(
        {"/data/": RAMResource()},
        mode=MountMode.WRITE,
        profiles=_role(commands=CommandsBlock(deny=(
            CommandRule(reason="prod is protected", paths=(
                "/data/prod/*", )), ))),
    )
    try:
        await ws.execute("mkdir -p /data/other")
        await ws.ops.write("/data/other/ok.txt", b"fine\n")
        with pytest.raises(PermissionError) as excinfo:
            await ws.ops.write("/data/prod/x.txt", b"nope\n")
        assert excinfo.value.errno == errno.EACCES
        assert "prod is protected" in str(excinfo.value)
        with pytest.raises(PermissionError):
            await ws.ops.read("/data/prod/x.txt")
    finally:
        await ws.close()


class SuppressProdWrites(Policy):

    async def post_ops(self, ctx: OpsResultContext) -> Action | None:
        if ctx.write and ctx.path.virtual.startswith("/data/prod/"):
            return Deny("write suppressed")
        return None


@pytest.mark.asyncio
async def test_touch_on_an_existing_file_is_a_write_at_the_op_door():
    # touch on an existing file mutates via setattr, not create; the
    # write classification must cover that op too.
    ws = Workspace({"/data/": RAMResource()}, mode=MountMode.WRITE)
    try:
        await ws.execute("mkdir -p /data/prod")
        await ws.ops.write("/data/prod/x.txt", b"keep\n")
        ws.policies.add(ReadOnlyProd())
        result = await ws.execute("touch /data/prod/x.txt")
        assert result.exit_code != 0
        assert b"Permission denied" in result.stderr
    finally:
        await ws.close()


@pytest.mark.asyncio
async def test_post_ops_deny_still_records_the_completed_write():
    # A post deny suppresses the result, not the effect: the backend
    # already mutated, so observation and caches must reflect the op.
    ws = Workspace({"/data/": RAMResource()}, mode=MountMode.WRITE)
    try:
        await ws.execute("mkdir -p /data/prod")
        ws.policies.add(SuppressProdWrites())
        with pytest.raises(PermissionError):
            await ws.ops.write("/data/prod/x.txt", b"data\n")
        assert any(r.op == "write" for r in ws.ops.records)
        assert await ws.ops.read("/data/prod/x.txt") == b"data\n"
    finally:
        await ws.close()


class SuppressProdReads(Policy):

    async def post_ops(self, ctx: OpsResultContext) -> Action | None:
        if not ctx.write and ctx.path.virtual.startswith("/data/prod/"):
            return Deny("no reads")
        return None


@pytest.mark.asyncio
async def test_post_ops_deny_records_the_bytes_a_denied_read_moved():
    # The suppressed result is the only place a read's byte count
    # lived, so without carrying it on the exception the record says
    # zero and network_bytes under-reports traffic that happened.
    ws = Workspace({"/data/": RAMResource()}, mode=MountMode.WRITE)
    try:
        await ws.execute("mkdir -p /data/prod")
        await ws.ops.write("/data/prod/x.txt", b"0123456789")
        ws.ops.records.clear()
        ws.policies.add(SuppressProdReads())
        with pytest.raises(PermissionError):
            await ws.ops.read("/data/prod/x.txt")
        reads = [r for r in ws.ops.records if r.op == "read"]
        assert len(reads) == 1
        assert reads[0].bytes == 10
    finally:
        await ws.close()


class CapProdReads(Policy):

    async def post_ops(self, ctx: OpsResultContext) -> Action | None:
        if not ctx.write and ctx.path.virtual.startswith("/data/prod/"):
            return Limit(max_bytes=3)
        return None


class CachingRAM(RAMResource):
    caches_reads = True
    name = "s3"


@pytest.mark.asyncio
async def test_a_capped_read_records_what_the_backend_moved():
    # A post_ops Limit truncates what the caller receives; the transfer
    # already happened, so recording the capped length would under-report
    # network_bytes by whatever the cap removed.
    ws = Workspace({"/data/": RAMResource()}, mode=MountMode.WRITE)
    try:
        await ws.execute("mkdir -p /data/prod")
        await ws.ops.write("/data/prod/x.txt", b"0123456789")
        ws.ops.records.clear()
        ws.policies.add(CapProdReads())
        assert await ws.ops.read("/data/prod/x.txt") == b"012"
        reads = [r for r in ws.ops.records if r.op == "read"]
        assert [r.bytes for r in reads] == [10]
    finally:
        await ws.close()


@pytest.mark.asyncio
async def test_a_denied_warm_read_is_not_counted_as_network_traffic():
    # The deny suppresses a result the cache produced, so nothing
    # crossed the network; recording it against the backend would count
    # traffic that never happened.
    ws = Workspace({"/data/": CachingRAM()}, mode=MountMode.WRITE)
    try:
        await ws.execute("mkdir -p /data/prod")
        await ws.ops.write("/data/prod/x.txt", b"0123456789")
        await ws.apply_io(
            IOResult(reads={"/data/prod/x.txt": b"0123456789"},
                     cache=["/data/prod/x.txt"]))
        ws.ops.records.clear()
        ws.policies.add(SuppressProdReads())
        with pytest.raises(PermissionError):
            await ws.ops.read("/data/prod/x.txt")
        rec = ws.ops.records[-1]
        assert rec.source == "ram"
        assert rec.is_cache is True
        assert ws.ops.network_bytes == 0
    finally:
        await ws.close()


class HardCapProdReads(Policy):

    async def post_ops(self, ctx: OpsResultContext) -> Action | None:
        if not ctx.write and ctx.path.virtual.startswith("/data/prod/"):
            return Limit(max_bytes=3, on_exceed=OnExceed.ERROR)
        return None


class ColdRemote(RAMResource):
    caches_reads = False
    name = "s3"


@pytest.mark.asyncio
async def test_a_hard_capped_read_records_what_the_backend_moved():
    # An ERROR-mode cap refuses the caller the bytes, but the backend
    # already moved them; dropping the record loses the whole transfer
    # rather than just truncating it.
    ws = Workspace({"/data/": ColdRemote()}, mode=MountMode.WRITE)
    try:
        await ws.execute("mkdir -p /data/prod")
        await ws.ops.write("/data/prod/x.txt", b"0123456789")
        ws.ops.records.clear()
        ws.policies.add(HardCapProdReads())
        with pytest.raises(LimitExceededError):
            await ws.ops.read("/data/prod/x.txt")
        reads = [r for r in ws.ops.records if r.op == "read"]
        assert [(r.source, r.bytes) for r in reads] == [("s3", 10)]
        assert ws.ops.network_bytes == 10
    finally:
        await ws.close()


@pytest.mark.asyncio
async def test_a_hard_capped_warm_read_is_not_network_traffic():
    # Same refusal, but the cache produced the bytes, so the transfer
    # it stands for never happened.
    ws = Workspace({"/data/": CachingRAM()}, mode=MountMode.WRITE)
    try:
        await ws.execute("mkdir -p /data/prod")
        await ws.ops.write("/data/prod/x.txt", b"0123456789")
        await ws.apply_io(
            IOResult(reads={"/data/prod/x.txt": b"0123456789"},
                     cache=["/data/prod/x.txt"]))
        ws.ops.records.clear()
        ws.policies.add(HardCapProdReads())
        with pytest.raises(LimitExceededError):
            await ws.ops.read("/data/prod/x.txt")
        rec = ws.ops.records[-1]
        assert rec.source == "ram"
        assert rec.is_cache is True
        assert ws.ops.network_bytes == 0
    finally:
        await ws.close()


class BrokenPostOps(Policy):

    async def post_ops(self, ctx: OpsResultContext):
        if ctx.write and ctx.path.virtual == "/data/prod/x.txt":
            return 42
        return None


@pytest.mark.asyncio
async def test_a_committed_write_is_recorded_when_bookkeeping_fails():
    # The backend applied the write, then a step after it (here an
    # invalid post_ops return, but any foreign bookkeeping error looks
    # the same) blew up. The error must propagate AND the transfer must
    # stay on the books: the door stamped the report at completion, so
    # the record does not depend on what kind of exception followed.
    ws = Workspace({"/data/": ColdRemote()}, mode=MountMode.WRITE)
    try:
        await ws.execute("mkdir -p /data/prod")
        ws.ops.records.clear()
        ws.policies.add(BrokenPostOps())
        with pytest.raises(PolicyError):
            await ws.ops.write("/data/prod/x.txt", b"123456")
        assert await ws.ops.read("/data/prod/x.txt") == b"123456"
        writes = [r for r in ws.ops.records if r.op == "write"]
        assert [(r.source, r.bytes) for r in writes] == [("s3", 6)]
        assert ws.ops.network_bytes >= 6
    finally:
        await ws.close()


@pytest.mark.asyncio
async def test_pre_ops_policy_holds_on_the_dispatcher_door():
    # touch routes through the shell's internal dispatcher, not
    # handle_command; a pre_ops-only policy must still refuse it.
    ws = Workspace({"/data/": RAMResource()}, mode=MountMode.WRITE)
    try:
        ws.policies.add(ReadOnlyProd())
        await ws.execute("mkdir -p /data/prod")
        result = await ws.execute("touch /data/prod/x")
        assert result.exit_code != 0
        assert b"Permission denied" in result.stderr
        ok = await ws.execute("touch /data/free && echo done")
        assert b"done" in ok.stdout
    finally:
        await ws.close()


class CapLines(Policy):

    async def post_execute(self, ctx: ExecuteResultContext) -> Action | None:
        return Limit(max_lines=2)


class CapReadBytes(Policy):

    async def post_ops(self, ctx: OpsResultContext) -> Action | None:
        if ctx.op == "read":
            return Limit(max_bytes=4)
        return None


@pytest.mark.asyncio
async def test_user_limit_policy_caps_line_output():
    # A user Limit merges with the built-in cap (tightest wins) and
    # bounds what execute() returns.
    ws = Workspace({"/data/": RAMResource()}, mode=MountMode.WRITE)
    try:
        ws.policies.add(CapLines())
        await ws.ops.write("/data/big.txt", b"1\n2\n3\n4\n5\n")
        r = await ws.execute("cat /data/big.txt")
        assert (await r.stdout_str()).count("\n") == 2
        assert "output truncated" in (await r.stderr_str())
    finally:
        await ws.close()


@pytest.mark.asyncio
async def test_user_limit_policy_caps_op_reads():
    # A post_ops Limit bounds the programmatic door too: ws.ops (and
    # FUSE behind it) serve capped bytes.
    ws = Workspace({"/data/": RAMResource()}, mode=MountMode.WRITE)
    try:
        ws.policies.add(CapReadBytes())
        await ws.ops.write("/data/f.txt", b"hello world")
        assert await ws.ops.read("/data/f.txt") == b"hell"
    finally:
        await ws.close()


class CapBytesHard(Policy):

    async def post_execute(self, ctx: ExecuteResultContext) -> Action | None:
        return Limit(max_bytes=4, on_exceed=OnExceed.ERROR)


class Boom(Policy):

    async def post_execute(self, ctx: ExecuteResultContext) -> Action | None:
        raise RuntimeError("boom")


class DenyReads(Policy):

    async def post_ops(self, ctx: OpsResultContext) -> Action | None:
        if ctx.op == "read":
            return Deny("reads are suppressed")
        return None


class SeeProducer(Policy):

    def __init__(self) -> None:
        self.seen: list[str] = []

    async def post_execute(self, ctx: ExecuteResultContext) -> Action | None:
        self.seen.append(ctx.producer.command)
        return None


@pytest.mark.asyncio
async def test_two_limit_policies_merge_to_the_tightest_end_to_end():
    ws = Workspace({"/data/": RAMResource()}, mode=MountMode.WRITE)
    try:
        ws.policies.add(CapLines())
        ws.policies.add(SuppressNothingCapThree())
        await ws.ops.write("/data/big.txt", b"1\n2\n3\n4\n5\n")
        r = await ws.execute("cat /data/big.txt")
        # CapLines says 2, SuppressNothingCapThree says 3: tightest wins.
        assert (await r.stdout_str()).count("\n") == 2
    finally:
        await ws.close()


class SuppressNothingCapThree(Policy):

    async def post_execute(self, ctx: ExecuteResultContext) -> Action | None:
        return Limit(max_lines=3)


@pytest.mark.asyncio
async def test_error_mode_limit_fails_the_line():
    # ANY-error: a user policy in error mode turns overflow into exit 1
    # with no stdout, GNU-style notice on stderr.
    ws = Workspace({"/data/": RAMResource()}, mode=MountMode.WRITE)
    try:
        ws.policies.add(CapBytesHard())
        await ws.ops.write("/data/f.txt", b"hello world\n")
        r = await ws.execute("cat /data/f.txt")
        assert r.exit_code == 1
        assert r.stdout is None or await r.stdout_str() == ""
        assert "output truncated" in (await r.stderr_str())
        ok = await ws.execute("echo ok")
        assert ok.exit_code == 0  # within the bound: no refusal
        assert await ok.stdout_str() == "ok\n"
    finally:
        await ws.close()


@pytest.mark.asyncio
async def test_a_post_ops_deny_beats_a_limit():
    # A refusal suppresses the result; bounding it would be meaningless.
    ws = Workspace({"/data/": RAMResource()}, mode=MountMode.WRITE)
    try:
        ws.policies.add(CapReadBytes())
        ws.policies.add(DenyReads())
        await ws.ops.write("/data/f.txt", b"hello world")
        with pytest.raises(PermissionError) as excinfo:
            await ws.ops.read("/data/f.txt")
        assert "reads are suppressed" in str(excinfo.value)
    finally:
        await ws.close()


@pytest.mark.asyncio
async def test_a_raising_post_execute_policy_fails_the_line_closed():
    ws = Workspace({"/data/": RAMResource()}, mode=MountMode.WRITE)
    try:
        ws.policies.add(Boom())
        r = await ws.execute("echo hi")
        assert r.exit_code == 126
        err = await r.stderr_str()
        assert err == "echo: policy denied: policy Boom failed: boom\n"
        assert r.stdout is None or await r.stdout_str() == ""
    finally:
        await ws.close()


@pytest.mark.asyncio
async def test_post_execute_sees_the_rightmost_producer():
    # The provenance a policy reads follows shell semantics: the tail
    # of a pipe, the right side of `;` and `||`.
    ws = Workspace({"/data/": RAMResource()}, mode=MountMode.WRITE)
    try:
        spy = SeeProducer()
        ws.policies.add(spy)
        await ws.ops.write("/data/f.txt", b"a\nb\n")
        await ws.execute("cat /data/f.txt | wc -l")
        await ws.execute("cat /data/f.txt ; head -n 1 /data/f.txt")
        await ws.execute("false || cat /data/f.txt")
        # Builtins carry provenance too: a policy keyed on echo sees it.
        await ws.execute("echo hi")
        await ws.execute("cat /data/f.txt ; echo done")
        assert spy.seen == ["wc", "head", "cat", "echo", "echo"]
    finally:
        await ws.close()


@pytest.mark.asyncio
async def test_role_hides_bind_every_session_including_the_default():
    ram = RAMResource()
    ws = Workspace({"/data/": ram},
                   mode=MountMode.WRITE,
                   profiles=_role(paths=PathsBlock(hide=("/data/finance",
                                                         "*.key"))))
    try:
        await ws.execute("mkdir -p /data/finance /data/pub")
        await ws.ops.write("/data/pub/a.txt", b"a\n")
        await ws.ops.write("/data/pub/b.key", b"k\n")
        # The default session cannot see the bound hides ...
        listing = await ws.execute("ls /data /data/pub")
        assert b"finance" not in listing.stdout
        assert b"b.key" not in listing.stdout
        assert b"a.txt" in listing.stdout
        # ... and neither can one created later from the same role.
        ws.create_session("late")
        gone = await ws.execute("cat /data/pub/b.key", session_id="late")
        assert gone.exit_code != 0
    finally:
        await ws.close()


@pytest.mark.asyncio
async def test_a_mount_sections_hides_are_written_in_full():
    # The section's entries are absolute, checked at load time to lie
    # under the mount root. They used to be relative and joined onto the
    # root, which silently doubled a path already written in full
    # (`/repo/secret` under `/repo` became `/repo/repo/secret`).
    repo = RAMResource()
    other = RAMResource()
    ws = Workspace(
        {
            "/repo/": (repo, MountMode.WRITE),
            "/other/": (other, MountMode.WRITE),
        },
        mode=MountMode.WRITE,
        profiles={
            "default": {
                "mounts": {
                    "/repo": {
                        "paths": {
                            "hide": ["/repo/.env", "/repo/*.pem"]
                        }
                    }
                }
            }
        })
    try:
        await ws.execute("mkdir -p /repo/certs /other")
        await ws.ops.write("/repo/.env", b"S=1\n")
        await ws.ops.write("/repo/certs/k.pem", b"pem\n")
        await ws.ops.write("/repo/README", b"r\n")
        await ws.ops.write("/other/.env", b"visible\n")
        await ws.ops.write("/other/x.pem", b"visible\n")
        listing = await ws.execute("ls -a /repo /repo/certs /other")
        out = listing.stdout.decode()
        # The section reaches only under its own root, so the same two
        # names outside the mount are untouched.
        assert ".env" in out.split("/other:")[1]
        assert "x.pem" in out
        assert "k.pem" not in out
        assert "README" in out
        repo_part = out.split("/other:")[0]
        assert ".env" not in repo_part.replace("/repo/certs:", "")
        hidden = await ws.execute("cat /repo/.env")
        assert hidden.exit_code != 0
        shown = await ws.execute("cat /other/.env")
        assert shown.stdout == b"visible\n"
    finally:
        await ws.close()


@pytest.mark.asyncio
async def test_a_bare_name_under_deny_refuses_with_the_default_reason():
    # The document's deny rules compile at construction; a bare string
    # is one command name with the default reason.
    ws = Workspace({"/data/": RAMResource()},
                   mode=MountMode.WRITE,
                   profiles=_role(commands=CommandsBlock(deny=("shred", ))))
    try:
        result = await ws.execute("shred /data/x")
        assert result.exit_code == 126
        assert result.stderr == b"shred: policy denied: denied by policy\n"
    finally:
        await ws.close()
