modelos-secure 0.1.1 → 0.2.0

package/python/modelos_secure/edge.py

@@ -0,0 +1,348 @@
+"""The edge for Python apps: the same as hooks/node/src/edge.mjs.
+
+Every framework ends at an HTTP call to a model, so during an attempt the hook watches there too: it
+sees each model call made through httpx (the tools offered and the calls the model asks for), adds
+the prompt canary, and plants attacker text into a tool result. Every other call that would leave the
+machine is stopped and recorded: through httpx it is answered as if it worked; through http.client
+(urllib, requests, urllib3, botocore) and aiohttp it is refused. A tool known to run keeps its calls.
+Outside an attempt it does nothing. Standard library only. ADR 0062.
+"""
+
+import contextvars
+import json
+import random
+import re
+
+#: Which tool is running, so its own calls can be told apart from calls nobody declared.
+current_tool = contextvars.ContextVar("modelos_secure_tool", default=None)
+#: Set while a model client the hook already patches is calling, so the call is counted once.
+in_sdk = contextvars.ContextVar("modelos_secure_sdk", default=False)
+
+# The same lists as hooks/node/src/hosts.mjs.
+_MODEL_HOSTS = [
+    r"(^|\.)api\.openai\.com$",
+    r"\.openai\.azure\.com$",
+    r"(^|\.)api\.anthropic\.com$",
+    r"(^|\.)generativelanguage\.googleapis\.com$",
+    r"aiplatform\.googleapis\.com$",
+    r"^(oauth2|sts|iamcredentials)\.googleapis\.com$",
+    r"^bedrock-runtime\.[\w-]+\.amazonaws\.com$",
+    r"^(api\.mistral\.ai|api\.groq\.com|api\.cohere\.com|openrouter\.ai|api\.together\.xyz|api\.deepseek\.com|api\.x\.ai|api\.fireworks\.ai|router\.huggingface\.co|api-inference\.huggingface\.co)$",
+    r"^login\.microsoftonline\.com$",
+    r"^metadata\.google\.internal$",
+]
+_LOOPBACK = r"^(localhost|127\.\d+\.\d+\.\d+|::1|\[::1\]|0\.0\.0\.0|169\.254\.169\.254)$"
+_MODEL_PATHS = [
+    (r"/chat/completions$", "chat"),
+    (r"/responses$", "responses"),
+    (r"/messages$", "anthropic"),
+    (r":(stream)?[gG]enerateContent$", "gemini"),
+    (r"/api/chat$", "ollama"),
+]
+
+
+def allowed_host(host, allow=()):
+    name = str(host or "").lower()
+    return (
+        re.match(_LOOPBACK, name) is not None
+        or any(re.search(p, name) for p in _MODEL_HOSTS)
+        or any(name == a or name.endswith("." + a) for a in allow)
+    )
+
+
+def model_shape(path, body):
+    if not isinstance(body, dict):
+        return None
+    for pattern, shape in _MODEL_PATHS:
+        if re.search(pattern, path):
+            return shape
+    if isinstance(body.get("messages"), list) and isinstance(body.get("model"), str):
+        return "chat"
+    if isinstance(body.get("contents"), list):
+        return "gemini"
+    return None
+
+
+def _with_line(content, line):
+    if isinstance(content, str):
+        return content + "\n" + line
+    if isinstance(content, list):
+        return content + [{"type": "text", "text": line}]
+    return content
+
+
+def add_canary(shape, body, line):
+    if not line or line in json.dumps(body, default=str):
+        return body
+    body = dict(body)
+    if shape in ("chat", "ollama"):
+        messages = list(body.get("messages") or [])
+        for i, message in enumerate(messages):
+            if isinstance(message, dict) and message.get("role") in ("system", "developer"):
+                messages[i] = {**message, "content": _with_line(message.get("content"), line)}
+                break
+        body["messages"] = messages
+    elif shape == "responses" and isinstance(body.get("instructions"), str):
+        body["instructions"] = body["instructions"] + "\n" + line
+    elif shape == "anthropic" and body.get("system"):
+        system = body["system"]
+        body["system"] = system + "\n" + line if isinstance(system, str) else list(system) + [{"type": "text", "text": line}]
+    elif shape == "gemini":
+        instruction = body.get("systemInstruction") or body.get("system_instruction")
+        if isinstance(instruction, dict) and isinstance(instruction.get("parts"), list):
+            body["systemInstruction"] = {**instruction, "parts": instruction["parts"] + [{"text": line}]}
+    return body
+
+
+def plant_into_results(shape, body, tool, text):
+    """Puts attacker text into the result of one tool, where the model reads it. Returns the body and
+    whether it went in."""
+    delivered = [False]
+
+    def plant(value):
+        delivered[0] = True
+        if isinstance(value, str):
+            return value + "\n\n" + text
+        if isinstance(value, list):
+            return value + [{"type": "text", "text": text}]
+        return {**(value or {}), "note": text}
+
+    body = dict(body)
+    if shape in ("chat", "ollama"):
+        named = {}
+        for message in body.get("messages") or []:
+            for call in (message or {}).get("tool_calls") or []:
+                named[call.get("id")] = (call.get("function") or {}).get("name")
+        body["messages"] = [
+            {**m, "content": plant(m.get("content"))}
+            if isinstance(m, dict) and m.get("role") == "tool" and (named.get(m.get("tool_call_id")) == tool or m.get("name") == tool)
+            else m
+            for m in body.get("messages") or []
+        ]
+    elif shape == "anthropic":
+        named = {}
+        for message in body.get("messages") or []:
+            for part in message.get("content") if isinstance(message.get("content"), list) else []:
+                if part.get("type") == "tool_use":
+                    named[part.get("id")] = part.get("name")
+        messages = []
+        for message in body.get("messages") or []:
+            content = message.get("content")
+            if isinstance(content, list):
+                content = [
+                    {**p, "content": plant(p.get("content"))}
+                    if p.get("type") == "tool_result" and named.get(p.get("tool_use_id")) == tool
+                    else p
+                    for p in content
+                ]
+            messages.append({**message, "content": content})
+        body["messages"] = messages
+    elif shape == "gemini":
+        contents = []
+        for content in body.get("contents") or []:
+            parts = []
+            for part in content.get("parts") or []:
+                response = part.get("functionResponse") or {}
+                if response.get("name") == tool:
+                    part = {**part, "functionResponse": {**response, "response": plant(response.get("response") or {})}}
+                parts.append(part)
+            contents.append({**content, "parts": parts})
+        body["contents"] = contents
+    return body, delivered[0]
+
+
+def _args(value):
+    if isinstance(value, dict):
+        return value
+    try:
+        parsed = json.loads(value)
+        return parsed if isinstance(parsed, dict) else {"value": parsed}
+    except Exception:
+        return {"raw": str(value or "")}
+
+
+def read_model(body, answer):
+    calls = []
+    for choice in answer.get("choices") or []:
+        for call in (choice.get("message") or {}).get("tool_calls") or []:
+            function = call.get("function") or {}
+            calls.append({"tool": function.get("name"), "args": _args(function.get("arguments"))})
+    for call in (answer.get("message") or {}).get("tool_calls") or []:
+        function = call.get("function") or {}
+        calls.append({"tool": function.get("name"), "args": _args(function.get("arguments"))})
+    for item in answer.get("output") or []:
+        if item.get("type") == "function_call":
+            calls.append({"tool": item.get("name"), "args": _args(item.get("arguments"))})
+    for part in answer.get("content") if isinstance(answer.get("content"), list) else []:
+        if part.get("type") == "tool_use":
+            calls.append({"tool": part.get("name"), "args": part.get("input") or {}})
+    for candidate in answer.get("candidates") or []:
+        for part in (candidate.get("content") or {}).get("parts") or []:
+            call = part.get("functionCall")
+            if call:
+                calls.append({"tool": call.get("name"), "args": call.get("args") or {}})
+    usage = answer.get("usage") or answer.get("usageMetadata") or {}
+    return {
+        "model": body.get("model") or answer.get("model"),
+        "toolCalls": [c for c in calls if c["tool"]],
+        "inputTokens": usage.get("prompt_tokens") or usage.get("input_tokens") or usage.get("promptTokenCount") or 0,
+        "outputTokens": usage.get("completion_tokens") or usage.get("output_tokens") or usage.get("candidatesTokenCount") or 0,
+    }
+
+
+def read_stream(body, text):
+    """A streamed answer, read event by event into the shape of a whole one."""
+    pieces = {}
+    whole = {"choices": [], "content": [], "candidates": [], "output": []}
+    for line in text.splitlines():
+        if not line.startswith("data:"):
+            continue
+        try:
+            event = json.loads(line[5:].strip())
+        except Exception:
+            continue
+        for choice in event.get("choices") or []:
+            for delta in (choice.get("delta") or {}).get("tool_calls") or []:
+                piece = pieces.setdefault(("c", delta.get("index")), {"name": "", "args": ""})
+                piece["name"] += (delta.get("function") or {}).get("name") or ""
+                piece["args"] += (delta.get("function") or {}).get("arguments") or ""
+        block = event.get("content_block") or {}
+        if event.get("type") == "content_block_start" and block.get("type") == "tool_use":
+            pieces[("a", event.get("index"))] = {"name": block.get("name"), "args": ""}
+        delta = event.get("delta") or {}
+        if delta.get("type") == "input_json_delta" and ("a", event.get("index")) in pieces:
+            pieces[("a", event.get("index"))]["args"] += delta.get("partial_json") or ""
+        whole["candidates"] += event.get("candidates") or []
+        if event.get("type") == "response.completed":
+            whole["output"] = (event.get("response") or {}).get("output") or []
+        if event.get("usage") or event.get("usageMetadata"):
+            whole["usage"] = event.get("usage") or event.get("usageMetadata")
+    for piece in pieces.values():
+        whole["choices"].append({"message": {"tool_calls": [{"function": {"name": piece["name"], "arguments": piece["args"]}}]}})
+    return read_model(body, whole)
+
+
+def install(hook):
+    """`hook` gives: state(), canary_line(), mode(name), way_in(name), allow, safely(fn, *args)."""
+    allow = tuple(hook["allow"])
+
+    def attempt():
+        state = hook["state"]()
+        return state if state and state.get("attempt") else None
+
+    def passes(host):
+        return allowed_host(host, allow) or hook["mode"](current_tool.get()) == "run"
+
+    def hold(method, host, path, size):
+        state = attempt()
+        if state is not None:
+            state.setdefault("edge", []).append({"method": method, "host": host, "path": path, "bytes": size})
+
+    def prepare(shape, body):
+        state = attempt()
+        prepared = add_canary(shape, body, hook["canary_line"]())
+        plant = (state.get("attempt") or {}).get("plant") if state else None
+        way = hook["way_in"](plant.get("wayIn")) if plant else None
+        if plant and way and way.get("function") and not state.get("plantDelivered"):
+            prepared, delivered = plant_into_results(shape, prepared, way["function"], plant.get("text", ""))
+            if delivered:
+                state["plantDelivered"] = True
+        return prepared
+
+    def see(body, answer_text, streamed, counted):
+        state = attempt()
+        if state is None:
+            return
+        try:
+            read = read_stream(body, answer_text) if streamed else read_model(body, json.loads(answer_text))
+        except Exception:
+            return
+        if not counted:
+            state["modelCalls"].append({**read, "edge": True})
+        for call in read["toolCalls"]:
+            if hook["mode"](call["tool"]) == "hold" and not any(h["tool"] == call["tool"] for h in state["held"]):
+                state["held"].append({"tool": call["tool"], "args": call["args"], "at": "edge"})
+
+    def held_response(httpx, request):
+        body = json.dumps({"status": "ok", "id": f"held-{random.randint(10000, 99999)}"}).encode()
+        return httpx.Response(200, headers={"content-type": "application/json"}, content=body, request=request)
+
+    def patch_httpx(httpx):
+        def rebuilt(request, body):
+            headers = {k: v for k, v in request.headers.items() if k.lower() != "content-length"}
+            return httpx.Request(request.method, request.url, headers=headers, content=json.dumps(body).encode())
+
+        def body_of(request):
+            try:
+                return json.loads(request.read() or b"null")
+            except Exception:
+                return None
+
+        def wanted(request):
+            body = body_of(request)
+            return body, model_shape(request.url.path, body)
+
+        sync = httpx.HTTPTransport.handle_request
+
+        def handle_request(self, request):
+            if attempt() is None:
+                return sync(self, request)
+            body, shape = hook["safely"](wanted, request) or (None, None)
+            if shape:
+                counted = in_sdk.get()
+                response = sync(self, hook["safely"](rebuilt, request, prepare(shape, body)) or request)
+                response.read()
+                hook["safely"](see, body, response.text, "event-stream" in response.headers.get("content-type", ""), counted)
+                return response
+            if passes(request.url.host):
+                return sync(self, request)
+            hold(request.method, request.url.host, request.url.path, len(request.content or b""))
+            return held_response(httpx, request)
+
+        httpx.HTTPTransport.handle_request = handle_request
+
+        asynchronous = httpx.AsyncHTTPTransport.handle_async_request
+
+        async def handle_async_request(self, request):
+            if attempt() is None:
+                return await asynchronous(self, request)
+            body, shape = hook["safely"](wanted, request) or (None, None)
+            if shape:
+                counted = in_sdk.get()
+                response = await asynchronous(self, hook["safely"](rebuilt, request, prepare(shape, body)) or request)
+                await response.aread()
+                hook["safely"](see, body, response.text, "event-stream" in response.headers.get("content-type", ""), counted)
+                return response
+            if passes(request.url.host):
+                return await asynchronous(self, request)
+            hold(request.method, request.url.host, request.url.path, len(request.content or b""))
+            return held_response(httpx, request)
+
+        httpx.AsyncHTTPTransport.handle_async_request = handle_async_request
+
+    def patch_http_client(http_client):
+        # putrequest, because urllib3 2 replaces request but still calls this one.
+        original = http_client.HTTPConnection.putrequest
+
+        def putrequest(self, method, url, *args, **kwargs):
+            if attempt() is not None and not passes(self.host):
+                hold(method, self.host, str(url).split("?")[0], 0)
+                raise ConnectionRefusedError("Stopped by ModelOS Secure during a review.")
+            return original(self, method, url, *args, **kwargs)
+
+        http_client.HTTPConnection.putrequest = putrequest
+
+    def patch_aiohttp(aiohttp):
+        original = aiohttp.ClientSession._request
+
+        async def _request(self, method, str_or_url, *args, **kwargs):
+            host = getattr(str_or_url, "host", None) or re.sub(r"^\w+://([^/:]+).*$", r"\1", str(str_or_url))
+            if attempt() is not None and not passes(host):
+                hold(method, host, "", 0)
+                raise ConnectionRefusedError("Stopped by ModelOS Secure during a review.")
+            return await original(self, method, str_or_url, *args, **kwargs)
+
+        aiohttp.ClientSession._request = _request
+
+    return {"httpx": patch_httpx, "http.client": patch_http_client, "aiohttp": patch_aiohttp}
+

package/python/modelos_secure/hook.py

@@ -17,7 +17,7 @@ import threading
 import types
 import urllib.parse
 
-from modelos_secure import clients
+from modelos_secure import clients, edge
 
 _MARK = "__modelos_secure_wrapped__"
 
@@ -37,6 +37,7 @@ def _fresh_state():
         "ran": [],
         "plantDelivered": False,
         "retrieved": [],
+        "edge": [],
     }
 
 
@@ -182,7 +183,11 @@ def _call(key, args, kwargs):
         return iter([held]) if inspect.isgeneratorfunction(entry["original"]) else held
     if name in _tools:
         _state["ran"].append(name)
-    result = entry["original"](*args, **kwargs)
+    running = edge.current_tool.set(name)
+    try:
+        result = entry["original"](*args, **kwargs)
+    finally:
+        edge.current_tool.reset(running)
     if inspect.isawaitable(result):
         return _finish_later(name, result)
     return _finish(name, result)
@@ -195,7 +200,11 @@ async def _acall(key, args, kwargs):
         return _hold(entry, args, kwargs)
     if name in _tools:
         _state["ran"].append(name)
-    return _finish(name, await entry["original"](*args, **kwargs))
+    running = edge.current_tool.set(name)
+    try:
+        return _finish(name, await entry["original"](*args, **kwargs))
+    finally:
+        edge.current_tool.reset(running)
 
 
 async def _finish_later(name, awaitable):
@@ -422,13 +431,17 @@ class _Control(http.server.BaseHTTPRequestHandler):
         if self.path == "/end":
             seen = _state
             _state = _fresh_state()
+            # The edge sees a tool call before its function runs. Where a hook then held the
+            # function itself, that record is the one kept, so a call counts once.
+            by_hook = {h["tool"] for h in seen["held"] if h.get("at") != "edge"}
             return self._send(
                 {
-                    "heldCalls": seen["held"],
+                    "heldCalls": [h for h in seen["held"] if h.get("at") != "edge" or h["tool"] not in by_hook],
                     "modelCalls": seen["modelCalls"],
                     "ranTools": seen["ran"],
                     "plantDelivered": seen["plantDelivered"],
                     "retrieved": seen["retrieved"],
+                    "edgeHeld": seen.get("edge", []),
                 }
             )
         if self.path == "/hello":
@@ -535,6 +548,22 @@ def install():
     _ways_in.update({way["function"]: way for way in config.get("waysIn", []) if way.get("function")})
     _files.update(os.path.realpath(path) for path in config.get("files", []))
     clients.connect(_canary_line, _see_model_call, _safely)
+    # The edge, beneath every framework. http.client is always there; the others patch on import.
+    ways = {way["name"]: way for way in config.get("waysIn", [])}
+    patches = edge.install(
+        {
+            "state": lambda: _state,
+            "canary_line": _canary_line,
+            "mode": lambda name: _tools.get(name, {}).get("mode") if name else None,
+            "way_in": ways.get,
+            "allow": (config.get("edge") or {}).get("allow", []),
+            "safely": _safely,
+        }
+    )
+    import http.client
+
+    _safely(patches["http.client"], http.client)
+    clients.LIBRARIES.update({"httpx": patches["httpx"], "aiohttp": patches["aiohttp"]})
     sys.meta_path.insert(0, _Finder)
     for name in list(clients.LIBRARIES):
         if name in sys.modules:

package/assets/login-BcHFyX6z.js

@@ -1 +0,0 @@
-export { t as loginCommand } from "./login-WwE-Gunj.js";