Update to v1.2.0

Author: pdudotdev

README.md

@@ -1,4 +1,4 @@
-# YANA - Yet Another Network Agent
+# YANA • Yet Another Network Agent
 
 [![Version](https://img.shields.io/badge/version-1.2-1a1a2e)](https://github.com/pdudotdev/YANA/releases/tag/v1.2.0)
 ![License](https://img.shields.io/badge/license-GPLv3-1a1a2e)
@@ -52,7 +52,7 @@ Run your tests with any framework. When something fails, YANA investigates - it
 
 **Step 1 - Install and ingest:**
 ```bash
-sudo apt install git make python3.12-venv
+sudo apt install git make python3.12-venv -y
 cd ~ && git clone https://github.com/pdudotdev/YANA
 cd YANA && make setup
 ```

input_models/models.py

@@ -1,15 +1,17 @@
 """Pydantic input models for MCP tools."""
+import ipaddress
 import json
 import re
 from typing import Literal, Optional
 
 from pydantic import BaseModel, Field, field_validator, model_validator
 
 _VRF_RE = re.compile(r'^[a-zA-Z0-9_-]{1,32}$')
+_DEVICE_RE = re.compile(r'^[A-Za-z0-9_-]{1,64}$')
 
 
 class BaseParamsModel(BaseModel):
-    """Base class with JSON string parsing and VRF validation."""
+    """Base class with JSON string parsing, device and VRF validation."""
 
     @model_validator(mode='before')
     @classmethod
@@ -22,6 +24,15 @@ def parse_string_input(cls, v):
                 raise ValueError(f"Could not parse params as JSON: {v!r}") from e
         return v
 
+    @field_validator('device', mode='before', check_fields=False)
+    @classmethod
+    def _validate_device(cls, v):
+        if v is None:
+            return v
+        if not _DEVICE_RE.match(str(v)):
+            raise ValueError(f"device must be alphanumeric with underscores/dashes, max 64 chars. Got: {v!r}")
+        return v
+
     @field_validator('vrf', mode='before', check_fields=False)
     @classmethod
     def _validate_vrf(cls, v):
@@ -64,10 +75,21 @@ class IntentQuery(BaseParamsModel):
 
 class TracerouteInput(BaseParamsModel):
     device: str = Field(..., description="Device name from inventory")
-    destination: str = Field(..., description="Destination IP address")
-    source: Optional[str] = Field(None, description="Source IP address (forces traceroute to use this interface)")
+    destination: str = Field(..., description="Destination IP address", max_length=45)
+    source: Optional[str] = Field(None, description="Source IP address (forces traceroute to use this interface)", max_length=45)
     vrf: str | None = Field(None, description="Optional VRF name")
 
+    @field_validator('destination', 'source', mode='before')
+    @classmethod
+    def _validate_ip(cls, v):
+        if v is None:
+            return v
+        try:
+            ipaddress.ip_address(str(v).strip())
+        except ValueError:
+            raise ValueError(f"Invalid IP address: {v!r}")
+        return str(v).strip()
+
 
 class KBQuery(BaseParamsModel):
     query: str = Field(..., description="Search question for the network knowledge base", max_length=500)

platforms/platform_map.py

@@ -124,7 +124,7 @@
         },
         "tools": {
             # count=1 limits to one probe per hop so the command terminates (default is continuous)
-            "traceroute": "/tool/traceroute count=1",
+            "traceroute": {"default": "/tool/traceroute count=1", "vrf": "/tool/traceroute count=1 routing-table={vrf}"},
         },
     },
 

testing/automated/test_suite.py

@@ -8,14 +8,15 @@
 from pydantic import ValidationError
 
 from input_models.models import (
-    DeviceListQuery, IntentQuery, InterfacesQuery, KBQuery,
-    OspfQuery, TracerouteInput,
+    DeviceListQuery, IntentQuery, KBQuery,
+    OspfQuery, RoutingQuery, TracerouteInput,
 )
-from platforms.platform_map import PLATFORM_MAP, _apply_vrf, get_action
+from platforms.platform_map import PLATFORM_MAP, _apply_vrf
 from tools.intent import query_intent
 from tools.inventory_tool import list_devices
 from tools.operational import traceroute
 from tools.ospf import get_ospf
+from tools.routing import get_routing
 from tools.status import get_status
 from transport.ssh import _build_cli, execute_ssh
 
@@ -177,6 +178,15 @@ async def test_traceroute_source_appended():
     assert "source 192.168.1.1" in result["_command"]
 
 
+async def test_routing_full_stack_ios():
+    with patch("transport.execute_ssh", new_callable=AsyncMock) as mock_ssh:
+        mock_ssh.return_value = "O    10.0.0.0/24 [110/20] via 10.0.0.1"
+        result = await get_routing(RoutingQuery(device="R1", query="ip_route"))
+    assert result["device"] == "R1"
+    assert result["cli_style"] == "ios"
+    assert result["_command"] == "show ip route vrf VRF1"
+
+
 # ── Inventory & device listing ───────────────────────────────────────────────
 
 def test_device_lookup_returns_dict(monkeypatch):
@@ -201,6 +211,10 @@ async def test_status_structure():
         mi.exists.return_value = False
         result = await get_status()
     assert set(result.keys()) == {"inventory", "intent", "chromadb"}
+    assert result["inventory"]["device_count"] == 6
+    assert result["inventory"]["source"] == "network_json"
+    assert result["intent"]["source"] == "unavailable"
+    assert result["chromadb"]["available"] is False
 
 
 async def test_intent_single_device_filter(tmp_path):
@@ -238,7 +252,22 @@ async def test_kb_basic_query():
 
 # ── Adversarial input ────────────────────────────────────────────────────────
 
-@pytest.mark.parametrize("bad_device", ["; rm -rf /", "' OR 1=1 --", "a" * 1000, ""])
-async def test_adversarial_device_name(bad_device):
-    result = await get_ospf(OspfQuery(device=bad_device, query="neighbors"))
-    assert "error" in result
+_BAD_DEVICES = [
+    "; rm -rf /", "' OR 1=1 --", "a" * 65, "",
+    "R1 extra", "R1\nshow run", "R1;cmd", "$(reboot)",
+    "R1\x00cmd", "R1&& cat /etc/shadow",
+]
+
+_VALID_DEVICES = ["R1", "my_device", "core-rtr_01", "A" * 64]
+
+
+@pytest.mark.parametrize("bad_device", _BAD_DEVICES)
+def test_adversarial_device_name_blocked(bad_device):
+    with pytest.raises(ValidationError):
+        OspfQuery(device=bad_device, query="neighbors")
+
+
+@pytest.mark.parametrize("good_device", _VALID_DEVICES)
+def test_valid_device_names_accepted(good_device):
+    q = OspfQuery(device=good_device, query="neighbors")
+    assert q.device == good_device

tools/operational.py

@@ -60,7 +60,9 @@ async def traceroute(params: TracerouteInput) -> dict:
         command = f"{base_cmd} {params.destination}"
         if params.source:
             command += f" source {params.source}"
-        if cli_style == "ios":
+        if cli_style in ("ios", "eos", "vyos"):
             command += " probe 1 timeout 2"
+        elif cli_style == "junos":
+            command += " wait 1"
 
     return await execute_command(params.device, command, timeout_ops=SSH_TIMEOUT_OPS_LONG)

tools/rag.py

@@ -1,6 +1,7 @@
 """RAG tool: search the network knowledge base."""
 import asyncio
 import logging
+import threading
 
 from input_models.models import KBQuery
 from tools import CHROMA_DIR
@@ -15,19 +16,22 @@
 # the device tools (get_ospf, get_interfaces) from loading.
 _embeddings = None
 _vectorstore = None
+_init_lock = threading.Lock()
 
 
 def _get_vectorstore():
     global _embeddings, _vectorstore
     if _vectorstore is None:
-        from langchain_huggingface import HuggingFaceEmbeddings
-        from langchain_chroma import Chroma
-        _embeddings = HuggingFaceEmbeddings(model_name=_EMBEDDING_MODEL)
-        _vectorstore = Chroma(
-            persist_directory=_CHROMA_DIR,
-            embedding_function=_embeddings,
-            collection_name=_COLLECTION,
-        )
+        with _init_lock:
+            if _vectorstore is None:
+                from langchain_huggingface import HuggingFaceEmbeddings
+                from langchain_chroma import Chroma
+                _embeddings = HuggingFaceEmbeddings(model_name=_EMBEDDING_MODEL)
+                _vectorstore = Chroma(
+                    persist_directory=_CHROMA_DIR,
+                    embedding_function=_embeddings,
+                    collection_name=_COLLECTION,
+                )
     return _vectorstore