aofire-python-agent 0.1.0__py3-none-any.whl

python_agent/dag_integrity.py

@@ -0,0 +1,198 @@
+"""HMAC integrity verification and injection scanning."""
+
+from __future__ import annotations
+
+import hashlib
+import hmac
+import json
+import os
+import re
+from typing import Any
+
+_INJECTION_PATTERNS: list[re.Pattern[str]] = [
+    re.compile(p, re.IGNORECASE) for p in [
+        r"ignore\s+(all\s+)?previous\s+instructions",
+        r"disregard\s+(all\s+)?previous",
+        r"you\s+are\s+now\s+a",
+        r"new\s+instructions:",
+        r"system\s+prompt:",
+        r"</ontology-data>",
+        r"</strategy-data>",
+        r"</candidate-summaries>",
+        r"</context-data>",
+        r"</user-input>",
+    ]
+]
+
+
+def generate_key() -> str:
+    """Generate 32-byte random key, hex-encoded."""
+    return os.urandom(32).hex()
+
+
+def load_or_create_key(path: str) -> str:
+    """Load hex key from file, or create file with new key."""
+    try:
+        with open(path) as f:
+            return f.read().strip()
+    except FileNotFoundError:
+        key = generate_key()
+        with open(path, "w") as f:
+            f.write(key)
+        return key
+
+
+def compute_hash(
+    ontology_dict: dict[str, Any], key: str,
+) -> str:
+    """HMAC-SHA256 hex digest of deterministic JSON."""
+    payload = json.dumps(ontology_dict, sort_keys=True)
+    return hmac.new(
+        bytes.fromhex(key),
+        payload.encode(),
+        hashlib.sha256,
+    ).hexdigest()
+
+
+def sign_node(node: Any, key: str) -> None:
+    """Set node.integrity_hash from ontology content."""
+    node.integrity_hash = compute_hash(
+        node.ontology.model_dump(), key,
+    )
+
+
+def verify_node(node: Any, key: str) -> bool:
+    """Return True if hash matches. False if tampered.
+
+    Returns False for unsigned nodes (empty hash).
+    """
+    if not node.integrity_hash:
+        return False
+    expected = compute_hash(
+        node.ontology.model_dump(), key,
+    )
+    return hmac.compare_digest(
+        node.integrity_hash, expected,
+    )
+
+
+def verify_dag(dag: Any, key: str) -> list[str]:
+    """Return IDs of signed nodes that fail verification.
+
+    Unsigned nodes (empty hash) are skipped.
+    """
+    failed: list[str] = []
+    for n in dag.nodes:
+        if not n.integrity_hash:
+            continue
+        if not verify_node(n, key):
+            failed.append(n.id)
+    return failed
+
+
+def scan_text_for_injection(text: str) -> list[str]:
+    """Scan a string for common injection patterns.
+
+    Returns list of matched pattern descriptions.
+    """
+    matches: list[str] = []
+    for pattern in _INJECTION_PATTERNS:
+        if pattern.search(text):
+            matches.append(pattern.pattern)
+    return matches
+
+
+def _collect_entity_texts(
+    ontology_dict: dict[str, Any],
+) -> list[str]:
+    """Extract text fields from entities."""
+    texts: list[str] = []
+    for entity in ontology_dict.get("entities", []):
+        texts.append(entity.get("description", ""))
+    return texts
+
+
+def _collect_relationship_texts(
+    ontology_dict: dict[str, Any],
+) -> list[str]:
+    """Extract text fields from relationships."""
+    texts: list[str] = []
+    for r in ontology_dict.get("relationships", []):
+        texts.append(r.get("description", ""))
+    return texts
+
+
+def _collect_constraint_texts(
+    ontology_dict: dict[str, Any],
+) -> list[str]:
+    """Extract text fields from domain constraints."""
+    texts: list[str] = []
+    for c in ontology_dict.get("domain_constraints", []):
+        texts.append(c.get("description", ""))
+        texts.append(c.get("expression", ""))
+    return texts
+
+
+def _collect_module_texts(
+    ontology_dict: dict[str, Any],
+) -> list[str]:
+    """Extract text fields from modules, classes, functions."""
+    texts: list[str] = []
+    for m in ontology_dict.get("modules", []):
+        texts.append(m.get("responsibility", ""))
+        texts.append(m.get("test_strategy", ""))
+        for cls in m.get("classes", []):
+            texts.append(cls.get("description", ""))
+            for fn in cls.get("methods", []):
+                texts.append(fn.get("docstring", ""))
+        for fn in m.get("functions", []):
+            texts.append(fn.get("docstring", ""))
+    return texts
+
+
+def _collect_misc_texts(
+    ontology_dict: dict[str, Any],
+) -> list[str]:
+    """Extract text from data_models, dependencies, questions."""
+    texts: list[str] = []
+    for dm in ontology_dict.get("data_models", []):
+        texts.append(dm.get("notes", ""))
+    for dep in ontology_dict.get("external_dependencies", []):
+        texts.append(dep.get("reason", ""))
+    for q in ontology_dict.get("open_questions", []):
+        texts.append(q.get("text", ""))
+        texts.append(q.get("context", ""))
+        texts.append(q.get("resolution", ""))
+    return texts
+
+
+def _collect_text_fields(
+    ontology_dict: dict[str, Any],
+) -> list[str]:
+    """Extract all free-text fields from an ontology dict."""
+    texts: list[str] = []
+    texts += _collect_entity_texts(ontology_dict)
+    texts += _collect_relationship_texts(ontology_dict)
+    texts += _collect_constraint_texts(ontology_dict)
+    texts += _collect_module_texts(ontology_dict)
+    texts += _collect_misc_texts(ontology_dict)
+    return [t for t in texts if t]
+
+
+def scan_ontology_for_injection(
+    ontology_dict: dict[str, Any],
+) -> list[str]:
+    """Scan all text fields in an ontology for injection.
+
+    Returns list of warnings (empty if clean).
+    """
+    warnings_list: list[str] = []
+    for text in _collect_text_fields(ontology_dict):
+        hits = scan_text_for_injection(text)
+        for pattern in hits:
+            preview = text[:80]
+            warnings_list.append(
+                f"Suspicious pattern {pattern!r} "
+                f"in: {preview!r}",
+            )
+    return warnings_list

python_agent/dag_utils.py

@@ -0,0 +1,181 @@
+"""Shared DAG persistence and snapshot utilities."""
+
+from __future__ import annotations
+
+import os
+import tempfile
+import uuid
+import warnings
+from datetime import datetime, timezone
+from pathlib import Path
+
+from python_agent.dag_integrity import (
+    load_or_create_key,
+    scan_ontology_for_injection,
+    sign_node,
+    verify_dag,
+)
+from python_agent.ontology import (
+    DAGEdge,
+    DAGNode,
+    Decision,
+    Ontology,
+    OntologyDAG,
+)
+
+
+def _default_key_path(dag_path: str) -> str:
+    """Derive key file path from DAG file path."""
+    return str(Path(dag_path).parent / ".dag-key")
+
+
+def _sign_unsigned_nodes(
+    dag: OntologyDAG, key: str,
+) -> None:
+    """Sign all nodes missing an integrity hash."""
+    for node in dag.nodes:
+        if not node.integrity_hash:
+            sign_node(node, key)
+
+
+def _verify_loaded_dag(
+    dag: OntologyDAG, key_path: str,
+) -> None:
+    """Verify all nodes and warn on failures."""
+    try:
+        key = load_or_create_key(key_path)
+    except (OSError, ValueError):
+        return
+    failed = verify_dag(dag, key)
+    if failed:
+        warnings.warn(
+            "DAG integrity check failed for nodes: "
+            + ", ".join(failed),
+            stacklevel=3,
+        )
+
+
+def _scan_loaded_dag(dag: OntologyDAG) -> None:
+    """Scan all nodes for injection patterns."""
+    for node in dag.nodes:
+        hits = scan_ontology_for_injection(
+            node.ontology.model_dump(),
+        )
+        for hit in hits:
+            warnings.warn(
+                f"Node {node.id}: {hit}",
+                stacklevel=3,
+            )
+
+
+def _read_file(path: str) -> str | None:
+    """Read file contents, or return None if not found."""
+    try:
+        with open(path) as f:
+            return f.read()
+    except FileNotFoundError:
+        return None
+
+
+def _parse_dag(text: str) -> OntologyDAG | None:
+    """Parse DAG JSON, or return None with warning."""
+    try:
+        return OntologyDAG.from_json(text)
+    except Exception as exc:
+        warnings.warn(
+            f"DAG validation error: {exc}",
+            stacklevel=3,
+        )
+        return None
+
+
+def load_dag(
+    path: str, project_name: str,
+    key_path: str | None = None,
+) -> OntologyDAG:
+    """Load an OntologyDAG from a JSON file.
+
+    Returns a new empty DAG if not found or invalid.
+    """
+    text = _read_file(path)
+    if text is None:
+        return OntologyDAG(project_name=project_name)
+    dag = _parse_dag(text)
+    if dag is None:
+        return OntologyDAG(project_name=project_name)
+    if key_path is None:
+        key_path = _default_key_path(path)
+    _verify_loaded_dag(dag, key_path)
+    _scan_loaded_dag(dag)
+    return dag
+
+
+def save_dag(
+    dag: OntologyDAG, path: str,
+    key_path: str | None = None,
+) -> None:
+    """Save an OntologyDAG, signing unsigned nodes.
+
+    Uses atomic write (temp file + rename) to prevent
+    corruption from interrupted writes.
+    """
+    if key_path is None:
+        key_path = _default_key_path(path)
+    key = load_or_create_key(key_path)
+    _sign_unsigned_nodes(dag, key)
+    parent_dir = os.path.dirname(os.path.abspath(path))
+    fd = tempfile.NamedTemporaryFile(
+        mode="w", dir=parent_dir,
+        suffix=".tmp", delete=False,
+    )
+    try:
+        fd.write(dag.to_json())
+        fd.close()
+        os.rename(fd.name, path)
+    except BaseException:
+        fd.close()
+        os.unlink(fd.name)
+        raise
+
+
+def make_node_id() -> str:
+    """Generate a unique node ID using uuid4."""
+    return str(uuid.uuid4())
+
+
+def save_snapshot(
+    dag: OntologyDAG, ontology: Ontology,
+    label: str, decision: Decision | None = None,
+) -> str:
+    """Create a new DAG node from the current ontology.
+
+    Links it as a child of the current node if one exists.
+    If decision is None, a default decision is used.
+    Returns the new node id.
+    """
+    now = datetime.now(timezone.utc).isoformat()
+    node_id = make_node_id()
+    node = DAGNode(
+        id=node_id,
+        ontology=ontology.model_copy(deep=True),
+        created_at=now,
+        label=label,
+    )
+    dag.nodes.append(node)
+    if dag.current_node_id:
+        if decision is None:
+            decision = Decision(
+                question="save",
+                options=["continue"],
+                chosen="continue",
+                rationale=label,
+            )
+        edge = DAGEdge(
+            parent_id=dag.current_node_id,
+            child_id=node_id,
+            decision=decision,
+            created_at=now,
+        )
+        dag.edges.append(edge)
+    dag.current_node_id = node_id
+    return node_id