Coverage for src/otg_mcp/client_capture.py: 8%
90 statements
« prev ^ index » next coverage.py v7.8.0, created at 2025-05-19 00:41 -0700
« prev ^ index » next coverage.py v7.8.0, created at 2025-05-19 00:41 -0700
1"""
2OTG Client Capture module providing specialized capture functionality.
4This module contains improved implementations for packet capture operations
5with proper control state handling.
6"""
8import logging
9import os
10import uuid
11from typing import Dict, List, Optional, Union, Any
13logger = logging.getLogger(__name__)
16def start_capture(api: Any, port_names: Union[str, List[str]]) -> Dict[str, Any]:
17 """
18 Start packet capture on one or more ports with proper control state handling.
20 Args:
21 api: Snappi API client
22 port_names: List or single name of port(s) to capture on
24 Returns:
25 Dictionary with status and result information
26 """
27 logger.info(f"Starting capture for ports: {port_names}")
29 logger.debug("Converting port names to list for consistent handling")
30 port_list = [port_names] if isinstance(port_names, str) else list(port_names)
32 try:
33 logger.debug("Creating control state with all required choices properly set")
34 cs = api.control_state()
36 logger.debug("Setting first-level choice to PORT")
37 cs.choice = cs.PORT
39 logger.debug("Setting second-level choice to CAPTURE for port control state")
40 cs.port.choice = cs.port.CAPTURE
42 logger.debug("Setting third-level choice: capture state to START")
43 cs.port.capture.state = cs.port.capture.START
45 logger.debug(f"Setting port names to capture on: {port_list}")
46 cs.port.capture.port_names = port_list
48 logger.debug("Applying control state")
49 logger.info(f"Setting control state to start capture on ports: {port_list}")
50 result = api.set_control_state(cs)
52 logger.debug("Checking for warnings in the result")
53 warnings = []
54 if hasattr(result, "warnings") and result.warnings:
55 warnings = result.warnings
56 logger.info(f"Start capture warnings: {warnings}")
58 return {"status": "success", "warnings": warnings}
60 except Exception as e:
61 logger.error(f"Error starting capture: {e}")
62 return {"status": "error", "error": str(e)}
65def stop_capture(api: Any, port_names: Union[str, List[str]]) -> Dict[str, Any]:
66 """
67 Stop packet capture on one or more ports with proper control state handling.
69 Args:
70 api: Snappi API client
71 port_names: List or single name of port(s) to stop capture on
73 Returns:
74 Dictionary with status and result information
75 """
76 logger.info(f"Stopping capture for ports: {port_names}")
78 logger.debug("Converting port names to list for consistent handling")
79 port_list = [port_names] if isinstance(port_names, str) else list(port_names)
81 try:
82 logger.debug("Creating control state with all required choices properly set")
83 cs = api.control_state()
85 logger.debug("Setting first-level choice to PORT")
86 cs.choice = cs.PORT
88 logger.debug("Setting second-level choice to CAPTURE for port control state")
89 cs.port.choice = cs.port.CAPTURE
91 logger.debug("Setting third-level choice: capture state to STOP")
92 cs.port.capture.state = cs.port.capture.STOP
94 logger.debug(f"Setting port names to capture on: {port_list}")
95 cs.port.capture.port_names = port_list
97 logger.debug("Applying control state")
98 logger.info(f"Setting control state to stop capture on ports: {port_list}")
99 result = api.set_control_state(cs)
101 logger.debug("Checking for warnings in the result")
102 warnings = []
103 if hasattr(result, "warnings") and result.warnings:
104 warnings = result.warnings
105 logger.info(f"Stop capture warnings: {warnings}")
107 return {"status": "success", "warnings": warnings}
109 except Exception as e:
110 logger.error(f"Error stopping capture: {e}")
111 return {"status": "error", "error": str(e)}
114def get_capture(
115 api: Any,
116 port_name: str,
117 output_dir: Optional[str] = None,
118 filename: Optional[str] = None,
119) -> Dict[str, Any]:
120 """
121 Get capture data from a port and save it to a file.
123 Args:
124 api: Snappi API client
125 port_name: Name of port to get capture from
126 output_dir: Directory to save the capture file (default: /tmp)
127 filename: Optional custom filename (default: auto-generated)
129 Returns:
130 Dictionary with status, file path, and capture data info
131 """
132 logger.info(f"Getting capture data for port {port_name}")
134 try:
135 logger.debug("Setting default output directory if not provided")
136 if output_dir is None:
137 output_dir = "/tmp"
139 logger.debug(f"Creating output directory if it doesn't exist: {output_dir}")
140 os.makedirs(output_dir, exist_ok=True)
142 logger.debug("Handling filename generation")
143 if filename is None:
144 filename = f"capture_{port_name}_{uuid.uuid4().hex[:8]}.pcap"
145 logger.debug(f"Generated unique filename: {filename}")
146 elif not filename.endswith(".pcap"):
147 logger.debug(f"Adding .pcap extension to filename: {filename}")
148 filename = f"{filename}.pcap"
150 file_path = os.path.join(output_dir, filename)
151 logger.info(f"Will save capture data to {file_path}")
153 logger.debug("Creating capture request with port name")
154 req = api.capture_request()
155 req.port_name = port_name
157 logger.debug("Requesting capture data from the device")
158 logger.info("Retrieving capture data")
159 pcap_bytes = api.get_capture(req)
161 logger.debug("Writing capture data to output file")
162 with open(file_path, "wb") as pcap_file:
163 pcap_file.write(pcap_bytes.read())
165 logger.info(f"Capture successfully saved to {file_path}")
167 return {
168 "status": "success",
169 "file_path": file_path,
170 "capture_id": filename,
171 "port": port_name,
172 "size_bytes": os.path.getsize(file_path),
173 }
175 except Exception as e:
176 logger.error(f"Error getting capture data: {e}")
177 return {
178 "status": "error",
179 "error": str(e),
180 "port": port_name,
181 "file_path": None,
182 "capture_id": None,
183 }