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

1""" 

2OTG Client Capture module providing specialized capture functionality. 

3 

4This module contains improved implementations for packet capture operations 

5with proper control state handling. 

6""" 

7 

8import logging 

9import os 

10import uuid 

11from typing import Dict, List, Optional, Union, Any 

12 

13logger = logging.getLogger(__name__) 

14 

15 

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. 

19 

20 Args: 

21 api: Snappi API client 

22 port_names: List or single name of port(s) to capture on 

23 

24 Returns: 

25 Dictionary with status and result information 

26 """ 

27 logger.info(f"Starting capture for ports: {port_names}") 

28 

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) 

31 

32 try: 

33 logger.debug("Creating control state with all required choices properly set") 

34 cs = api.control_state() 

35 

36 logger.debug("Setting first-level choice to PORT") 

37 cs.choice = cs.PORT 

38 

39 logger.debug("Setting second-level choice to CAPTURE for port control state") 

40 cs.port.choice = cs.port.CAPTURE 

41 

42 logger.debug("Setting third-level choice: capture state to START") 

43 cs.port.capture.state = cs.port.capture.START 

44 

45 logger.debug(f"Setting port names to capture on: {port_list}") 

46 cs.port.capture.port_names = port_list 

47 

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) 

51 

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}") 

57 

58 return {"status": "success", "warnings": warnings} 

59 

60 except Exception as e: 

61 logger.error(f"Error starting capture: {e}") 

62 return {"status": "error", "error": str(e)} 

63 

64 

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. 

68 

69 Args: 

70 api: Snappi API client 

71 port_names: List or single name of port(s) to stop capture on 

72 

73 Returns: 

74 Dictionary with status and result information 

75 """ 

76 logger.info(f"Stopping capture for ports: {port_names}") 

77 

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) 

80 

81 try: 

82 logger.debug("Creating control state with all required choices properly set") 

83 cs = api.control_state() 

84 

85 logger.debug("Setting first-level choice to PORT") 

86 cs.choice = cs.PORT 

87 

88 logger.debug("Setting second-level choice to CAPTURE for port control state") 

89 cs.port.choice = cs.port.CAPTURE 

90 

91 logger.debug("Setting third-level choice: capture state to STOP") 

92 cs.port.capture.state = cs.port.capture.STOP 

93 

94 logger.debug(f"Setting port names to capture on: {port_list}") 

95 cs.port.capture.port_names = port_list 

96 

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) 

100 

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}") 

106 

107 return {"status": "success", "warnings": warnings} 

108 

109 except Exception as e: 

110 logger.error(f"Error stopping capture: {e}") 

111 return {"status": "error", "error": str(e)} 

112 

113 

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. 

122 

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) 

128 

129 Returns: 

130 Dictionary with status, file path, and capture data info 

131 """ 

132 logger.info(f"Getting capture data for port {port_name}") 

133 

134 try: 

135 logger.debug("Setting default output directory if not provided") 

136 if output_dir is None: 

137 output_dir = "/tmp" 

138 

139 logger.debug(f"Creating output directory if it doesn't exist: {output_dir}") 

140 os.makedirs(output_dir, exist_ok=True) 

141 

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" 

149 

150 file_path = os.path.join(output_dir, filename) 

151 logger.info(f"Will save capture data to {file_path}") 

152 

153 logger.debug("Creating capture request with port name") 

154 req = api.capture_request() 

155 req.port_name = port_name 

156 

157 logger.debug("Requesting capture data from the device") 

158 logger.info("Retrieving capture data") 

159 pcap_bytes = api.get_capture(req) 

160 

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()) 

164 

165 logger.info(f"Capture successfully saved to {file_path}") 

166 

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 } 

174 

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 }