Coverage for src / server_list / spec / ups_collector.py: 88%

1#!/usr/bin/env python3 

2""" 

3NUT (Network UPS Tools) data collector. 

4 

5Communicates with NUT server via TCP socket to collect UPS information. 

6""" 

7 

8import logging 

9import socket 

10 

11import server_list.spec.models as models 

12 

13DEFAULT_PORT = 3493 

14SOCKET_TIMEOUT = 10 

15 

16 

17def _send_command(sock: socket.socket, command: str) -> list[str]: 

18 """Send a command to NUT server and receive response. 

19 

20 Args: 

21 sock: Connected socket 

22 command: NUT command to send 

23 

24 Returns: 

25 List of response lines 

26 """ 

27 sock.sendall(f"{command}\n".encode("utf-8")) 

28 

29 response = b"" 

30 while True: 

31 chunk = sock.recv(4096) 

32 if not chunk: 32 ↛ 33line 32 didn't jump to line 33 because the condition on line 32 was never true

33 break 

34 response += chunk 

35 # Check for end of response 

36 decoded = response.decode("utf-8") 

37 if "END LIST" in decoded or decoded.startswith("ERR "): 37 ↛ 30line 37 didn't jump to line 30 because the condition on line 37 was always true

38 break 

39 

40 return response.decode("utf-8").strip().split("\n") 

41 

42 

43def _parse_list_ups(lines: list[str]) -> list[tuple[str, str]]: 

44 """Parse LIST UPS response. 

45 

46 Args: 

47 lines: Response lines from LIST UPS command 

48 

49 Returns: 

50 List of (ups_name, description) tuples 

51 """ 

52 ups_list = [] 

53 for line in lines: 

54 if line.startswith("UPS "): 

55 # Format: UPS <upsname> "<description>" 

56 parts = line.split(" ", 2) 

57 if len(parts) >= 2: 57 ↛ 53line 57 didn't jump to line 53 because the condition on line 57 was always true

58 ups_name = parts[1] 

59 description = parts[2].strip('"') if len(parts) > 2 else "" 

60 ups_list.append((ups_name, description)) 

61 return ups_list 

62 

63 

64def _parse_list_var(lines: list[str]) -> dict[str, str]: 

65 """Parse LIST VAR response. 

66 

67 Args: 

68 lines: Response lines from LIST VAR command 

69 

70 Returns: 

71 Dict mapping variable name to value 

72 """ 

73 variables = {} 

74 for line in lines: 

75 if line.startswith("VAR "): 

76 # Format: VAR <upsname> <varname> "<value>" 

77 parts = line.split(" ", 3) 

78 if len(parts) >= 4: 78 ↛ 74line 78 didn't jump to line 74 because the condition on line 78 was always true

79 var_name = parts[2] 

80 value = parts[3].strip('"') 

81 variables[var_name] = value 

82 return variables 

83 

84 

85def _parse_list_client(lines: list[str]) -> list[str]: 

86 """Parse LIST CLIENT response. 

87 

88 Args: 

89 lines: Response lines from LIST CLIENT command 

90 

91 Returns: 

92 List of client IP addresses 

93 """ 

94 clients = [] 

95 for line in lines: 

96 if line.startswith("CLIENT "): 

97 # Format: CLIENT <upsname> <client_ip> 

98 parts = line.split(" ") 

99 if len(parts) >= 3: 99 ↛ 95line 99 didn't jump to line 95 because the condition on line 99 was always true

100 clients.append(parts[2]) 

101 return clients 

102 

103 

104def _safe_float(value: str | None) -> float | None: 

105 """Convert string to float safely.""" 

106 if value is None: 

107 return None 

108 try: 

109 return float(value) 

110 except (ValueError, TypeError): 

111 return None 

112 

113 

114def _safe_int(value: str | None) -> int | None: 

115 """Convert string to int safely.""" 

116 if value is None: 

117 return None 

118 try: 

119 return int(float(value)) 

120 except (ValueError, TypeError): 

121 return None 

122 

123 

124def connect_to_nut(host: str, port: int = DEFAULT_PORT) -> socket.socket | None: 

125 """Connect to NUT server. 

126 

127 Args: 

128 host: NUT server hostname 

129 port: NUT server port (default: 3493) 

130 

131 Returns: 

132 Connected socket or None if connection failed 

133 """ 

134 try: 

135 sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 

136 sock.settimeout(SOCKET_TIMEOUT) 

137 sock.connect((host, port)) 

138 return sock 

139 except (OSError, socket.error) as e: 

140 logging.warning("Failed to connect to NUT server %s:%d: %s", host, port, e) 

141 return None 

142 

143 

144def list_ups(sock: socket.socket) -> list[tuple[str, str]]: 

145 """List all UPS devices on the NUT server. 

146 

147 Args: 

148 sock: Connected socket 

149 

150 Returns: 

151 List of (ups_name, description) tuples 

152 """ 

153 try: 

154 lines = _send_command(sock, "LIST UPS") 

155 return _parse_list_ups(lines) 

156 except (OSError, socket.error) as e: 

157 logging.warning("Failed to list UPS: %s", e) 

158 return [] 

159 

160 

161def get_ups_variables(sock: socket.socket, ups_name: str) -> dict[str, str]: 

162 """Get all variables for a UPS. 

163 

164 Args: 

165 sock: Connected socket 

166 ups_name: UPS name 

167 

168 Returns: 

169 Dict mapping variable name to value 

170 """ 

171 try: 

172 lines = _send_command(sock, f"LIST VAR {ups_name}") 

173 return _parse_list_var(lines) 

174 except (OSError, socket.error) as e: 

175 logging.warning("Failed to get UPS variables for %s: %s", ups_name, e) 

176 return {} 

177 

178 

179def get_ups_clients(sock: socket.socket, ups_name: str) -> list[str]: 

180 """Get all clients connected to a UPS. 

181 

182 Args: 

183 sock: Connected socket 

184 ups_name: UPS name 

185 

186 Returns: 

187 List of client IP addresses 

188 """ 

189 try: 

190 lines = _send_command(sock, f"LIST CLIENT {ups_name}") 

191 return _parse_list_client(lines) 

192 except (OSError, socket.error) as e: 

193 logging.warning("Failed to get UPS clients for %s: %s", ups_name, e) 

194 return [] 

195 

196 

197def fetch_ups_info(host: str, ups_name: str, port: int = DEFAULT_PORT) -> models.UPSInfo | None: 

198 """Fetch UPS information from NUT server. 

199 

200 Args: 

201 host: NUT server hostname 

202 ups_name: UPS name 

203 port: NUT server port 

204 

205 Returns: 

206 UPSInfo or None if failed 

207 """ 

208 sock = connect_to_nut(host, port) 

209 if not sock: 

210 return None 

211 

212 try: 

213 variables = get_ups_variables(sock, ups_name) 

214 if not variables: 214 ↛ 215line 214 didn't jump to line 215 because the condition on line 214 was never true

215 return None 

216 

217 return models.UPSInfo( 

218 ups_name=ups_name, 

219 host=host, 

220 model=variables.get("ups.model"), 

221 battery_charge=_safe_float(variables.get("battery.charge")), 

222 battery_runtime=_safe_int(variables.get("battery.runtime")), 

223 ups_load=_safe_float(variables.get("ups.load")), 

224 ups_status=variables.get("ups.status"), 

225 ups_temperature=_safe_float(variables.get("ups.temperature")), 

226 input_voltage=_safe_float(variables.get("input.voltage")), 

227 output_voltage=_safe_float(variables.get("output.voltage")), 

228 ) 

229 finally: 

230 sock.close() 

231 

232 

233def fetch_ups_clients(host: str, ups_name: str, port: int = DEFAULT_PORT) -> list[models.UPSClient]: 

234 """Fetch UPS client information from NUT server. 

235 

236 Args: 

237 host: NUT server hostname 

238 ups_name: UPS name 

239 port: NUT server port 

240 

241 Returns: 

242 List of UPSClient objects 

243 """ 

244 sock = connect_to_nut(host, port) 

245 if not sock: 

246 return [] 

247 

248 try: 

249 client_ips = get_ups_clients(sock, ups_name) 

250 return [ 

251 models.UPSClient( 

252 ups_name=ups_name, 

253 host=host, 

254 client_ip=ip, 

255 client_hostname=None, # Hostname resolution can be added if needed 

256 ) 

257 for ip in client_ips 

258 ] 

259 finally: 

260 sock.close() 

261 

262 

263def fetch_all_ups_from_host( 

264 host: str, port: int = DEFAULT_PORT, ups_name_filter: str | None = None 

265) -> tuple[list[models.UPSInfo], list[models.UPSClient]]: 

266 """Fetch all UPS information from a NUT host. 

267 

268 Args: 

269 host: NUT server hostname 

270 port: NUT server port 

271 ups_name_filter: If specified, only fetch this UPS 

272 

273 Returns: 

274 Tuple of (list of UPSInfo, list of UPSClient) 

275 """ 

276 sock = connect_to_nut(host, port) 

277 if not sock: 

278 return [], [] 

279 

280 try: 

281 # Get list of UPS devices 

282 if ups_name_filter: 282 ↛ 283line 282 didn't jump to line 283 because the condition on line 282 was never true

283 ups_list = [(ups_name_filter, "")] 

284 else: 

285 ups_list = list_ups(sock) 

286 

287 all_ups_info: list[models.UPSInfo] = [] 

288 all_clients: list[models.UPSClient] = [] 

289 

290 for ups_name, _ in ups_list: 

291 # Get UPS variables 

292 variables = get_ups_variables(sock, ups_name) 

293 if variables: 293 ↛ 309line 293 didn't jump to line 309 because the condition on line 293 was always true

294 ups_info = models.UPSInfo( 

295 ups_name=ups_name, 

296 host=host, 

297 model=variables.get("ups.model"), 

298 battery_charge=_safe_float(variables.get("battery.charge")), 

299 battery_runtime=_safe_int(variables.get("battery.runtime")), 

300 ups_load=_safe_float(variables.get("ups.load")), 

301 ups_status=variables.get("ups.status"), 

302 ups_temperature=_safe_float(variables.get("ups.temperature")), 

303 input_voltage=_safe_float(variables.get("input.voltage")), 

304 output_voltage=_safe_float(variables.get("output.voltage")), 

305 ) 

306 all_ups_info.append(ups_info) 

307 

308 # Get UPS clients 

309 client_ips = get_ups_clients(sock, ups_name) 

310 for ip in client_ips: 

311 client = models.UPSClient( 

312 ups_name=ups_name, 

313 host=host, 

314 client_ip=ip, 

315 client_hostname=None, 

316 ) 

317 all_clients.append(client) 

318 

319 return all_ups_info, all_clients 

320 

321 finally: 

322 sock.close()