Coverage for src/otg_mcp/client_capture.py: 8%

1"""

2 OTG Client Capture module providing specialized capture functionality.

3

4 This module contains improved implementations for packet capture operations

5 with proper control state handling.

6"""

7

8 import logging

9 import os

10 import uuid

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

12

13 logger = logging.getLogger(__name__)

14

15

16 def 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

65 def 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

114 def 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 }