1#!/usr/bin/env python3

2

3 import argparse

4 import os

5 import subprocess as sp

6 import re

7 import shutil

8 import filecmp

9 import sys

10 import typing as T

11 from checkbox_support.scripts.image_checker import has_desktop_environment

12 from datetime import datetime

13

14

15# Checkbox could run in a snap container, so we need to prepend this root path

16 RUNTIME_ROOT = os.getenv("CHECKBOX_RUNTIME", default="")

17# Snap mount point, see

18# https://snapcraft.io/docs/environment-variables#heading--snap

19 SNAP = os.getenv("SNAP", default="")

20

21

22 def get_timestamp_str() -> str:

23 with open("/proc/uptime", "r") as f:

24 # uptime file always have 2 numbers

25 # uptime_seconds total_idle_seconds

26 # take the 1st one

27 uptime_seconds = f.readline().split()[0]

28

29 return "Time: {}; Uptime: {} seconds ".format(

30 datetime.now().strftime("%m/%d/%Y, %H:%M:%S"), uptime_seconds

31 )

32

33

34 class DeviceInfoCollector:

35

36 class Device:

37 PCI = "pci"

38 WIRELESS = "wireless"

39 USB = "usb"

40 DRM = "drm"

41

42 DEFAULT_DEVICES = {

43 "required": [

44 Device.WIRELESS,

45 Device.PCI,

46 Device.USB,

47 ], # these can fail the test case

48 "optional": [Device.DRM], # these only produce warnings

49 } # type: dict[str, list[str]]

50 # to modify, add more values in the enum

51 # and reference them in required/optional respectively

52

53 def get_drm_info(self) -> str:

54 return str(os.listdir("/sys/class/drm"))

55

56 def get_wireless_info(self) -> str:

57 iw_out = sp.check_output(["iw", "dev"], universal_newlines=True)

58 lines = iw_out.splitlines()

59 lines_to_write = list(

60 filter(

61 lambda line: "addr" in line

62 or "Interface" in line

63 or "ssid" in line,

64 sorted(lines),

65 )

66 )

67 return "\n".join(map(lambda line: line.strip(), lines_to_write))

68

69 def get_usb_info(self) -> str:

70 out = sp.check_output(

71 [

72 "checkbox-support-lsusb",

73 "-f",

74 '"{}"/var/lib/usbutils/usb.ids'.format(RUNTIME_ROOT),

75 "-s",

76 ],

77 universal_newlines=True,

78 ).splitlines()

79 out.sort()

80 return "\n".join(out)

81

82 def get_pci_info(self) -> str:

83 return sp.check_output(

84 ["lspci", "-i", "{}/usr/share/misc/pci.ids".format(SNAP)],

85 universal_newlines=True,

86 )

87

88 def compare_device_lists(

89 self,

90 expected_dir: str,

91 actual_dir: str,

92 devices: T.Dict[str, T.List[str]] = DEFAULT_DEVICES,

93 ) -> bool:

94 """Compares the list of devices in expected_dir against actual_dir

95

96 :param expected_dir: files containing the expected device list

97 :param actual_dir: files containing the actual device list

98 :param devices: what devices do we want to compare, see DEFAULT_DEVICES

99 :return: whether the device list matches

100 """

101 print(

102 "Comparing devices in (expected) {} against (actual) {}...".format(

103 expected_dir, actual_dir

104 )

105 )

106 for device in devices["required"]:

107 # file paths of the expected and actual device lists

108 expected = "{}/{}_log".format(expected_dir, device)

109 actual = "{}/{}_log".format(actual_dir, device)

110 if not filecmp.cmp(expected, actual):

111 print(

112 "[ ERR ] The output of {} differs!".format(device),

113 file=sys.stderr,

114 )

115 self.print_diff(device, expected, actual)

116 return False

117

118 for device in devices["optional"]:

119 expected = "{}/{}_log".format(expected_dir, device)

120 actual = "{}/{}_log".format(actual_dir, device)

121 if not filecmp.cmp(expected, actual):

122 print(

123 "[ WARN ] Items under {} have changed.".format(actual),

124 file=sys.stderr,

125 )

126 self.print_diff(device, expected, actual)

127

128 return True

129

130 def dump(

131 self,

132 output_directory: str,

133 devices: T.Dict[str, T.List[str]] = DEFAULT_DEVICES,

134 ) -> None:

135 os.makedirs(output_directory, exist_ok=True)

136 # add extra behavior if necessary

137 for device in devices["required"]:

138 with open(

139 "{}/{}_log".format(output_directory, device), "w"

140 ) as file:

141 file.write(self.dump_function[device]())

142

143 for device in devices["optional"]:

144 with open(

145 "{}/{}_log".format(output_directory, device), "w"

146 ) as file:

147 file.write(self.dump_function[device]())

148

149 os.sync()

150

151 def print_diff(self, name: str, expected_path: str, actual_path: str):

152 with open(expected_path) as file_expected, open(

153 actual_path

154 ) as file_actual:

155 print("Expected {} output:".format(name), file=sys.stderr)

156 print(file_expected.read(), file=sys.stderr)

157 print("Actual {} output:".format(name), file=sys.stderr)

158 print(file_actual.read(), file=sys.stderr)

159 print("End of {} diff".format(name), file=sys.stderr)

160

161 def __init__(self) -> None:

162 self.dump_function = {

163 self.Device.PCI: self.get_pci_info,

164 self.Device.DRM: self.get_drm_info,

165 self.Device.USB: self.get_usb_info,

166 self.Device.WIRELESS: self.get_wireless_info,

167 }

168

169

170 class FwtsTester:

171 def is_fwts_supported(self) -> bool:

172 return shutil.which("fwts") is not None

173

174 def fwts_log_check_passed(

175 self, output_directory: str, fwts_arguments=["klog", "oops"]

176 ) -> bool:

177 """

178 Check if fwts logs passes the checks specified in sleep_test_log_check

179 This script live in the same directory

180

181 :param output_directory: where the output of fwts should be written to

182 :type output_directory: str

183 :return: whether sleep_test_log_check.py returned 0 (success)

184 :rtype: bool

185 """

186 log_file_path = "{}/fwts_{}.log".format(

187 output_directory, "_".join(fwts_arguments)

188 )

189 sp.run(["fwts", "-r", log_file_path, "-q", *fwts_arguments])

190 result = sp.run(

191 [

192 "sleep_test_log_check.py",

193 "-v",

194 "--ignore-warning",

195 "-t",

196 "all",

197 log_file_path,

198 ],

199 )

200

201 return result.returncode == 0

202

203

204 class HardwareRendererTester:

205

206 def has_display_connection(self) -> bool:

207 """

208 Checks if a display is connected by searching /sys/class/drm

209

210 :return: True if there's at least 1 node that is connected

211 """

212

213 # look for GPU file nodes first

214 DRM_PATH = "/sys/class/drm"

215 possible_gpu_nodes = [

216 directory

217 for directory in os.listdir(DRM_PATH)

218 if directory != "version"

219 ]

220 if len(possible_gpu_nodes) == 0:

221 # kernel doesn't see any GPU nodes

222 print(

223 "There's nothing under {}".format(DRM_PATH),

224 "if an external GPU is connected,"

225 "check if the connection is loose.",

226 )

227 return False

228

229 print("Listing all DRM connection statuses:")

230

231 connected_to_display = False

232 for gpu in possible_gpu_nodes:

233 # for each gpu, check for connection

234 # return true if anything is connected

235 try:

236 with open("{}/{}/status".format(DRM_PATH, gpu)) as status_file:

237 status_str = status_file.read().strip().lower()

238 # - card0: connected

239 print(" - {}: {}".format(gpu, status_str))

240

241 if status_str == "connected":

242 connected_to_display = True

243 except FileNotFoundError:

244 # this just means we don't have a status file

245 # => no connection, continue to the next

246 pass

247 except Exception as e:

248 print("Unexpected error: ", e, file=sys.stderr)

249

250 if not connected_to_display:

251 print(

252 "No display is connected. This case will be skipped.",

253 "Maybe the display cable is not connected?",

254 "If the device is not supposed to have a display,"

255 "then skipping is expected.",

256 )

257

258 return connected_to_display

259

260 def is_hardware_renderer_available(self) -> bool:

261 """

262 Checks if hardware rendering is being used.

263 THIS ASSUMES A DRM CONNECTION EXISTS

264 - self.has_display_connection() should be called first if unsure

265

266 :return: True if a hardware renderer is active, otherwise return False

267 :rtype: bool

268 """

269

270 DISPLAY = os.getenv("DISPLAY", "")

271

272 if DISPLAY == "": 272 ↛ 273 line 272 didn't jump to line 273 because the condition on line 272 was never true

273 print("$DISPLAY is not set, we will let unity_support infer this")

274 else:

275 print("Checking $DISPLAY={}".format(DISPLAY))

276

277 unity_support_output = sp.run(

278 ["{}/usr/lib/nux/unity_support_test".format(RUNTIME_ROOT), "-p"],

279 stdout=sp.PIPE,

280 universal_newlines=True,

281 )

282 if unity_support_output.returncode != 0:

283 print(

284 "[ ERR ] unity support test returned {}. Error is: {}".format(

285 unity_support_output.returncode,

286 unity_support_output.stdout,

287 ),

288 file=sys.stderr,

289 )

290 return False

291

292 is_hardware_rendered = (

293 self.parse_unity_support_output(unity_support_output.stdout).get(

294 "Not software rendered"

295 )

296 == "yes"

297 )

298 if is_hardware_rendered:

299 print("[ OK ] This machine is using a hardware renderer!")

300 return True

301

302 print("[ ERR ] Software rendering detected", file=sys.stderr)

303 return False

304

305 def parse_unity_support_output(

306 self, unity_output_string: str

307 ) -> T.Dict[str, str]:

308 """

309 Parses the output of `unity_support_test` into a dictionary

310

311 :param output_string: the raw output from running unity_support_test -p

312 :type output_string: str

313 :return: string key-value pairs that mirror the output of unity_support

314 Left hand side of the first colon are the keys;

315 right hand side are the values.

316 :rtype: dict[str, str]

317 """

318

319 output = {} # type: dict[str, str]

320 for line in unity_output_string.split("\n"):

321 # max_split=1 to prevent splitting the string after the 1st colon

322 words = line.split(":", maxsplit=1)

323 if len(words) == 2:

324 key = words[0].strip()

325 value = remove_color_code(words[1].strip())

326 output[key] = value

327

328 return output

329

330

331 def get_failed_services() -> T.List[str]:

332 """

333 Counts the number of failed services listed in systemctl

334

335 :return: a list of failed services as they appear in systemctl

336 """

337 command = [

338 "systemctl",

339 "list-units",

340 "--system",

341 "--no-ask-password",

342 "--no-pager",

343 "--no-legend",

344 "--state=failed",

345 "--plain", # plaintext, otherwise it includes color codes

346 ]

347

348 return sp.check_output(command, universal_newlines=True).splitlines()

349

350

351 def create_parser():

352 parser = argparse.ArgumentParser(

353 prog="Reboot tests",

354 description="Collects device info and compares them across reboots",

355 )

356 parser.add_argument(

357 "-d",

358 "--dump-to",

359 required=False,

360 dest="output_directory",

361 help="Device info-dumps will be written here",

362 )

363 parser.add_argument(

364 "-c",

365 "--compare-to",

366 dest="comparison_directory",

367 help="Directory of ground-truth for device info comparison",

368 )

369 parser.add_argument(

370 "-s",

371 "--service-check",

372 default=False,

373 dest="do_service_check",

374 action="store_true",

375 help="If specified, check if all system services are running",

376 )

377 parser.add_argument(

378 "-f",

379 "--fwts-check",

380 default=False,

381 dest="do_fwts_check",

382 action="store_true",

383 help="If specified, look for fwts log errors",

384 )

385 parser.add_argument(

386 "-g",

387 "--graphics",

388 default=False,

389 dest="do_renderer_check",

390 action="store_true",

391 help="If specified, check if hardware rendering is being used",

392 )

393

394 return parser

395

396

397 def remove_color_code(string: str) -> str:

398 """

399 Removes ANSI color escape sequences from string

400

401 :param string: the string that you would like to remove color code

402 credit: Hanhsuan Lee <hanhsuan.lee@canonical.com>

403 """

404 return re.sub(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])", "", string)

405

406

407 def main() -> int:

408 """Main routine

409

410 :return: an return code for checkbox to consume, 1 = failed, 0 = success

411 :rtype: int

412 """

413

414 args = create_parser().parse_args()

415

416 # all 4 tests pass by default

417 # they only fail if their respective flags are specified

418 # if no flags are specified, calling this script is a no-op

419 fwts_passed = True

420 device_comparison_passed = True

421 renderer_test_passed = True

422 service_check_passed = True

423

424 print("Starting reboot checks. {}".format(get_timestamp_str()))

425

426 if args.comparison_directory is not None:

427 if args.output_directory is None:

428 print(

429 "[ ERR ] Please specify an output directory with the -d flag.",

430 file=sys.stderr,

431 )

432 raise ValueError(

433 "Comparison directory is specified, but output directory isn't"

434 )

435 else:

436 collector = DeviceInfoCollector()

437 collector.dump(args.output_directory)

438 if collector.compare_device_lists(

439 args.comparison_directory, args.output_directory

440 ):

441 print("[ OK ] Devices match!")

442 else:

443 device_comparison_passed = False

444

445 # dump (no checks) if only output_directory is specified

446 if args.output_directory is not None and args.comparison_directory is None:

447 print("Only dumping device info to {}".format(args.output_directory))

448 DeviceInfoCollector().dump(args.output_directory)

449

450 if args.do_fwts_check:

451 tester = FwtsTester()

452 if tester.is_fwts_supported() and not tester.fwts_log_check_passed( 452 ↛ 455 line 452 didn't jump to line 455 because the condition on line 452 was never true

453 args.output_directory

454 ):

455 fwts_passed = False

456 else:

457 print("[ OK ] fwts checks passed!")

458

459 if args.do_service_check:

460 failed_services = get_failed_services()

461 if len(failed_services) > 0:

462 print(

463 "These services failed: {}".format("\n".join(failed_services)),

464 file=sys.stderr,

465 )

466 service_check_passed = False

467 else:

468 print("[ OK ] Didn't find any failed system services!")

469

470 if args.do_renderer_check:

471 tester = HardwareRendererTester()

472 if has_desktop_environment() and tester.has_display_connection(): 472 ↛ 474 line 472 didn't jump to line 474 because the condition on line 472 was never true

473 # skip renderer test if there's no display

474 renderer_test_passed = tester.is_hardware_renderer_available()

475

476 print("Finished reboot checks. {}".format(get_timestamp_str()))

477

478 if (

479 fwts_passed

480 and device_comparison_passed

481 and renderer_test_passed

482 and service_check_passed

483 ):

484 return 0

485 else:

486 return 1

487

488

489 if __name__ == "__main__":

490 return_code = main()

491 exit(return_code)