feat: Add batch link_ids to link_capture_start/stop

This commit is contained in:
YueGuobin 2026-06-14 21:32:47 +08:00
parent e02f1a8cd0
commit 6df9374a4c
No known key found for this signature in database
2 changed files with 60 additions and 16 deletions

View File

@ -919,28 +919,34 @@ async def link_reset(
@mcp.tool()
async def link_capture_start(
project_id: Annotated[str, Field(description="UUID of the project")],
link_id: Annotated[str, Field(description="UUID of the link")],
link_id: Annotated[str | None, Field(description="Link UUID (single mode)")] = None,
data_link_type: Annotated[str, Field(description="Data link type (default: DLT_EN10MB)")] = "DLT_EN10MB",
capture_file_name: Annotated[str | None, Field(description="Capture file name (optional)")] = None,
wireshark: Annotated[bool, Field(description="Open Wireshark automatically (default: false)")] = False,
link_ids: Annotated[list[str] | None, Field(description="Batch mode: [\"uuid1\",\"uuid2\"] — start capture on multiple links in parallel")] = None,
) -> list[dict[str, Any]]:
"""Start packet capture on a link. The capture file can later be downloaded with download_capture_file."""
return await asyncio.to_thread(_run_handler_sync, start_capture_handler, {
"project_id": project_id, "link_id": link_id,
"data_link_type": data_link_type, "capture_file_name": capture_file_name,
"wireshark": wireshark,
})
"""Start packet capture on one or more links."""
params = {"project_id": project_id, "data_link_type": data_link_type, "capture_file_name": capture_file_name, "wireshark": wireshark}
if link_ids:
params["link_ids"] = link_ids
else:
params["link_id"] = link_id
return await asyncio.to_thread(_run_handler_sync, start_capture_handler, params)
@mcp.tool()
async def link_capture_stop(
project_id: Annotated[str, Field(description="UUID of the project")],
link_id: Annotated[str, Field(description="UUID of the link")],
link_id: Annotated[str | None, Field(description="Link UUID (single mode)")] = None,
link_ids: Annotated[list[str] | None, Field(description="Batch mode: [\"uuid1\",\"uuid2\"] — stop capture on multiple links in parallel")] = None,
) -> list[dict[str, Any]]:
"""Stop packet capture on a link. After stopping, the capture file can be downloaded."""
return await asyncio.to_thread(_run_handler_sync, stop_capture_handler, {
"project_id": project_id, "link_id": link_id,
})
"""Stop packet capture on one or more links."""
params = {"project_id": project_id}
if link_ids:
params["link_ids"] = link_ids
else:
params["link_id"] = link_id
return await asyncio.to_thread(_run_handler_sync, stop_capture_handler, params)
@mcp.tool()

View File

@ -157,11 +157,41 @@ def reset_link_handler(params: dict[str, Any], gns3_ctx: dict[str, Any]) -> dict
return {"message": f"Link {link_id} reset", "link": result}
def _batch_capture(project_id, link_ids, action, data_builder, conn):
"""Helper for batch capture start/stop."""
def _act(lid):
try:
url = f"{conn.base_url}/projects/{project_id}/links/{lid}/capture/{action}"
kwargs = data_builder(lid) if data_builder else {}
conn.http_call("post", url, **kwargs)
return {"link_id": lid, "status": "success"}
except Exception as e:
return {"link_id": lid, "status": "error", "error": str(e)}
with ThreadPoolExecutor(max_workers=min(len(link_ids), BATCH_MAX_WORKERS)) as pool:
return list(pool.map(_act, link_ids))
def start_capture_handler(params: dict[str, Any], gns3_ctx: dict[str, Any]) -> dict[str, Any]:
project_id = params.get("project_id")
if not project_id:
return {"error": "project_id is required"}
link_ids = params.get("link_ids")
if link_ids:
if not isinstance(link_ids, list):
return {"error": "link_ids must be a list"}
conn = _get_connector(gns3_ctx)
dlt = params.get("data_link_type", "DLT_EN10MB")
ws = params.get("wireshark", False)
fname = params.get("capture_file_name")
def _build(lid):
data = {"data_link_type": dlt, "wireshark": ws}
if fname:
data["capture_file_name"] = fname
return {"json_data": data}
return _batch_capture(project_id, link_ids, "start", _build, conn)
link_id = params.get("link_id")
if not project_id or not link_id:
return {"error": "project_id and link_id are required"}
if not link_id:
return {"error": "link_id or link_ids is required"}
conn = _get_connector(gns3_ctx)
data = {
"data_link_type": params.get("data_link_type", "DLT_EN10MB"),
@ -176,9 +206,17 @@ def start_capture_handler(params: dict[str, Any], gns3_ctx: dict[str, Any]) -> d
def stop_capture_handler(params: dict[str, Any], gns3_ctx: dict[str, Any]) -> dict[str, Any]:
project_id = params.get("project_id")
if not project_id:
return {"error": "project_id is required"}
link_ids = params.get("link_ids")
if link_ids:
if not isinstance(link_ids, list):
return {"error": "link_ids must be a list"}
conn = _get_connector(gns3_ctx)
return _batch_capture(project_id, link_ids, "stop", None, conn)
link_id = params.get("link_id")
if not project_id or not link_id:
return {"error": "project_id and link_id are required"}
if not link_id:
return {"error": "link_id or link_ids is required"}
conn = _get_connector(gns3_ctx)
url = f"{conn.base_url}/projects/{project_id}/links/{link_id}/capture/stop"
conn.http_call("post", url)