1#!/usr/bin/env python3

2"""

3 If you're debugging this program, set PM_TEST_DRY_RUN in your environment.

4 It will make the script not run actual S3, S4, reboot and poweroff commands.

5"""

6 import gi

7 import json

8 import logging

9 import logging.handlers

10 import os

11 import pwd

12 import re

13 import shutil

14 import signal

15 import subprocess

16 import sys

17 from argparse import ArgumentParser, SUPPRESS

18 from calendar import timegm

19 from configparser import ConfigParser

20 from datetime import datetime, timedelta

21 from time import localtime, time

22

23 gi.require_version("Gtk", "3.0")

24 from gi.repository import GObject, Gtk # noqa: E402

25

26

27 def main():

28 """

29 Run power management operation as many times as needed

30 """

31 args, extra_args = MyArgumentParser().parse()

32

33 # Verify that script is run as root

34 if os.getuid():

35 sys.stderr.write(

36 "This script needs superuser " "permissions to run correctly\n"

37 )

38 sys.exit(1)

39

40 # Obtain name of the invoking user.

41 username = os.getenv("NORMAL_USER")

42 if not username:

43 uid = os.getenv("SUDO_UID") or os.getenv("PKEXEC_UID")

44 if not uid:

45 sys.stderr.write("Unable to determine invoking user\n")

46 sys.exit(1)

47 username = pwd.getpwuid(int(uid)).pw_name

48

49 LoggingConfiguration.set(args.log_level, args.log_filename, args.append)

50 logging.debug("Invoking username: %s", username)

51 logging.debug("Arguments: {0!r}".format(args))

52 logging.debug("Extra Arguments: {0!r}".format(extra_args))

53

54 dry_run = os.environ.get("PM_TEST_DRY_RUN", False)

55 if dry_run:

56 logging.info("Running in dry-run mode")

57

58 try:

59 operation = PowerManagementOperation(

60 args, extra_args, user=username, dry_run=dry_run

61 )

62 operation.setup()

63 operation.run()

64 except (TestCancelled, TestFailed) as exception:

65 if isinstance(exception, TestFailed):

66 logging.error(exception.args[0])

67 message = exception.MESSAGE.format(args.pm_operation.capitalize())

68 if args.silent:

69 logging.info(message)

70 else:

71 title = "{0} test".format(args.pm_operation.capitalize())

72 MessageDialog(title, message, Gtk.MessageType.ERROR).run()

73 operation.teardown()

74 result = {

75 "outcome": "fail",

76 "comments": message,

77 }

78 with open(os.path.join(args.log_dir, "__result"), "wt") as f:

79 json.dump(result, f)

80 env = os.environ.copy()

81 # remove following envvars

82 for key in ["LD_LIBRARY_PATH", "PYTHONPATH", "PYTHONHOME"]:

83 if key in env.keys():

84 del env[key]

85 env["DISPLAY"] = ":0"

86 sudo_cmd = 'sudo -u {} bash -c "source {}; exec bash"'.format(

87 operation.user, args.checkbox_respawn_cmd

88 )

89 args = ["x-terminal-emulator", "-e", sudo_cmd]

90 print("\nCheckbox will resume in another window.")

91 print("It's safe to close this one.", flush=True)

92 os.execvpe("x-terminal-emulator", args, env)

93

94 return 0

95

96

97 class PowerManagementOperation:

98 SLEEP_TIME = 5

99

100 def __init__(self, args, extra_args, user=None, dry_run=False):

101 self.args = args

102 self.extra_args = extra_args

103 self.user = user

104 self.dry_run = dry_run

105 self.hw_list_start = os.path.join(

106 self.args.log_dir, "hardware.at_start"

107 )

108

109 def setup(self):

110 """

111 Enable configuration file

112 """

113 if self.args.check_hardware_list:

114 if not os.path.exists(self.hw_list_start):

115 # create baseline list once per session

116 with open(self.hw_list_start, "wt") as f:

117 f.write(self.get_hw_list())

118

119 # Enable autologin and sudo on first cycle

120 if self.args.total == self.args.repetitions:

121 AutoLoginConfigurator(user=self.user).enable()

122 SudoersConfigurator(user=self.user).enable()

123

124 # Schedule this script to be automatically executed

125 # on startup to continue testing

126 autostart_file = AutoStartFile(self.args, user=self.user)

127 autostart_file.write()

128

129 def run(self):

130 """

131 Run a power management iteration

132 """

133 logging.info(

134 "{0} operations remaining: {1}".format(

135 self.args.pm_operation, self.args.repetitions

136 )

137 )

138 if self.args.pm_timestamp:

139 pm_timestamp = datetime.fromtimestamp(self.args.pm_timestamp)

140 now = datetime.now()

141 pm_time = now - pm_timestamp

142 logging.info(

143 "{0} time: {1}".format(

144 self.args.pm_operation.capitalize(), pm_time

145 )

146 )

147 if self.args.repetitions > 0:

148 self.run_suspend_cycles(

149 self.args.suspends_before_reboot, self.args.fwts

150 )

151 self.run_pm_command()

152 if self.args.check_hardware_list:

153 self.check_hw_list()

154 else:

155 self.summary()

156

157 def run_pm_command(self):

158 """

159 Run power managment command and check result if needed

160 """

161 # Display information to user

162 # and make it possible to cancel the test

163 CountdownDialog(

164 self.args.pm_operation,

165 self.args.pm_delay,

166 self.args.hardware_delay,

167 self.args.total - self.args.repetitions,

168 self.args.total,

169 ).run()

170

171 # A small sleep time is added to reboot and poweroff

172 # so that script has time to return a value

173 # (useful when running it as an automated test)

174 command_str = "sleep {0}; {1}".format(

175 self.SLEEP_TIME, self.args.pm_operation

176 )

177 if self.extra_args:

178 command_str += " {0}".format(" ".join(self.extra_args))

179

180 if self.args.pm_operation != "reboot":

181 WakeUpAlarm.set(seconds=self.args.wakeup)

182

183 logging.info(

184 "Executing new {0!r} operation...".format(self.args.pm_operation)

185 )

186 logging.debug("Executing: {0!r}...".format(command_str))

187 if self.dry_run:

188 print("\n\nRUNNING IN DRY-RUN MODE")

189 print("Normally the program would run: {}".format(command_str))

190 print("Waiting for Enter instead")

191 input()

192 else:

193 # The PM operation is performed asynchronously so let's just wait

194 # indefinitely until it happens and someone comes along to kill us.

195 # This addresses LP: #1413134

196 subprocess.check_call(command_str, shell=True)

197 signal.pause()

198

199 def run_suspend_cycles(self, cycles_count, fwts):

200 """Run suspend and resume cycles."""

201 if cycles_count < 1:

202 return

203 script_path = ""

204 if fwts:

205 script_path = "checkbox-support-fwts_test"

206 command_tpl = (

207 "-s s3 --s3-device-check "

208 "--s3-sleep-delay=30 --s3-multiple={}"

209 )

210 if self.args.log_dir:

211 command_tpl = (

212 "--log={}/fwts.log ".format(self.args.log_dir)

213 + command_tpl

214 )

215 command_tpl = "{} " + command_tpl

216 else:

217 script_name = "sleep_test.py"

218 command_tpl = "{} -s mem -p -i {} -w 10"

219 script_path = os.path.join(

220 os.path.dirname(os.path.realpath(__file__)), script_name

221 )

222 command_str = command_tpl.format(script_path, cycles_count)

223 logging.info("Running suspend/resume cycles")

224 logging.debug("Executing: {0!r}...".format(command_str))

225 if self.dry_run:

226 print("\n\nRUNNING IN DRY-RUN MODE")

227 print("Normally the program would run: {}".format(command_str))

228 print("Waiting for Enter instead")

229 input()

230 else:

231 try:

232 # We call sleep_test.py or fwts_test script and log its output

233 # as it contains average times we need to compute global

234 # average times later.

235 logging.info(

236 subprocess.check_output(

237 command_str, universal_newlines=True, shell=True

238 )

239 )

240 except subprocess.CalledProcessError as e:

241 logging.error("Error while running {0}:".format(e.cmd))

242 logging.error(e.output)

243

244 def summary(self):

245 """

246 Gather hardware information for the last time,

247 log execution time and exit

248 """

249 # Just gather hardware information one more time and exit

250 CountdownDialog(

251 self.args.pm_operation,

252 self.args.pm_delay,

253 self.args.hardware_delay,

254 self.args.total - self.args.repetitions,

255 self.args.total,

256 ).run()

257

258 self.teardown()

259

260 # Log some time information

261 start = datetime.fromtimestamp(self.args.start)

262 end = datetime.now()

263 if self.args.pm_operation == "reboot":

264 sleep_time = timedelta(seconds=self.SLEEP_TIME)

265 else:

266 sleep_time = timedelta(seconds=self.args.wakeup)

267

268 wait_time = timedelta(

269 seconds=(

270 self.args.pm_delay + self.args.hardware_delay * self.args.total

271 )

272 )

273 average = (end - start - wait_time) / self.args.total - sleep_time

274 time_message = (

275 "Total elapsed time: {total}\n"

276 "Average recovery time: {average}".format(

277 total=end - start, average=average

278 )

279 )

280 logging.info(time_message)

281

282 message = "{0} test complete".format(

283 self.args.pm_operation.capitalize()

284 )

285 if self.args.suspends_before_reboot:

286 total_suspends_expected = (

287 self.args.suspends_before_reboot * self.args.total

288 )

289 problems = ""

290 fwts_log_path = os.path.join(self.args.log_dir, "fwts.log")

291 try:

292 with open(fwts_log_path, "rt") as fwts_log:

293 magic_line_s3 = "Completed S3 cycle(s) \n"

294 magic_line_s2idle = "Completed s2idle cycle(s) \n"

295 lines = fwts_log.readlines()

296 count_s3 = lines.count(magic_line_s3)

297 count_s2idle = lines.count(magic_line_s2idle)

298 count = count_s3 + count_s2idle

299 if count != total_suspends_expected:

300 problems = (

301 "Found {} occurrences of S3/S2idle."

302 " Expected {}".format(

303 count, total_suspends_expected

304 )

305 )

306 except FileNotFoundError:

307 problems = "Error opening {}".format(fwts_log_path)

308 if problems:

309 result = {

310 "outcome": "fail" if problems else "pass",

311 "comments": problems,

312 }

313 result_filename = os.path.join(self.args.log_dir, "__result")

314 with open(result_filename, "wt") as result_f:

315 json.dump(result, result_f)

316

317 if self.args.silent:

318 logging.info(message)

319 else:

320 title = "{0} test".format(self.args.pm_operation.capitalize())

321 MessageDialog(title, message).run()

322 if self.args.checkbox_respawn_cmd:

323 try:

324 subprocess.run(

325 r"unset LD_LIBRARY_PATH;"

326 r"unset PYTHONPATH; unset PYTHONHOME;"

327 r'DISPLAY=:0 x-terminal-emulator -e "sudo -u '

328 r'{} bash -c \"source {}; exec bash\""'.format(

329 self.user, self.args.checkbox_respawn_cmd

330 ),

331 shell=True,

332 check=True,

333 )

334 # x-terminal-emulator command does not work on Wayland

335 # Run the checkbox_respawn_cmd via subprocess.run instead

336 except subprocess.CalledProcessError:

337 with open(self.args.checkbox_respawn_cmd, "rt") as respawn_f:

338 for respawn_cmd in respawn_f:

339 subprocess.run(respawn_cmd, shell=True)

340

341 def teardown(self):

342 """

343 Restore configuration

344 """

345 # Don't execute this script again on next reboot

346 autostart_file = AutoStartFile(self.args, user=self.user)

347 autostart_file.remove()

348

349 # Restore previous configuration

350 SudoersConfigurator().disable()

351 AutoLoginConfigurator().disable()

352

353 def get_hw_list(self):

354 try:

355 content = subprocess.check_output(

356 "lspci", encoding=sys.stdout.encoding

357 )

358 content += subprocess.check_output(

359 "lsusb", encoding=sys.stdout.encoding

360 )

361 return content

362 except subprocess.CalledProcessError as exc:

363 logging.warning("Problem running lspci or lsusb: %s", exc)

364 return ""

365

366 def check_hw_list(self):

367 with open(self.hw_list_start, "rt") as f:

368 before = set(f.read().split("\n"))

369 after = set(self.get_hw_list().split("\n"))

370 if after != before:

371 message = "Hardware changed!"

372 only_before = before - after

373 if only_before:

374 message += "\nHardware lost after pm operation:"

375 for item in sorted(list(only_before)):

376 message += "\n\t{}".format(item)

377 only_after = after - before

378 if only_after:

379 message += "\nNew hardware found after pm operation:"

380 for item in sorted(list(only_after)):

381 message += "\n\t{}".format(item)

382 raise TestFailed(message)

383

384

385 class TestCancelled(Exception):

386 RETURN_CODE = 1

387 MESSAGE = "{0} test cancelled by user"

388

389

390 class TestFailed(Exception):

391 RETURN_CODE = 2

392 MESSAGE = "{0} test failed"

393

394

395 class WakeUpAlarm:

396 ALARM_FILENAME = "/sys/class/rtc/rtc0/wakealarm"

397 RTC_FILENAME = "/proc/driver/rtc"

398

399 @classmethod

400 def set(cls, minutes=0, seconds=0):

401 """

402 Calculate wakeup time and write it to BIOS

403 """

404 now = int(time())

405 timeout = minutes * 60 + seconds

406 wakeup_time_utc = now + timeout

407 wakeup_time_local = timegm(localtime()) + timeout

408

409 subprocess.check_call("echo 0 > %s" % cls.ALARM_FILENAME, shell=True)

410 subprocess.check_call(

411 "echo %d > %s" % (wakeup_time_utc, cls.ALARM_FILENAME), shell=True

412 )

413

414 with open(cls.ALARM_FILENAME) as alarm_file:

415 wakeup_time_stored_str = alarm_file.read()

416

417 if not re.match(r"\d+", wakeup_time_stored_str):

418 subprocess.check_call(

419 'echo "+%d" > %s' % (timeout, cls.ALARM_FILENAME),

420 shell=True,

421 )

422 with open(cls.ALARM_FILENAME) as alarm_file2:

423 wakeup_time_stored_str = alarm_file2.read()

424 if not re.match(r"\d+", wakeup_time_stored_str):

425 logging.error(

426 "Invalid wakeup time format: {0!r}".format(

427 wakeup_time_stored_str

428 )

429 )

430 sys.exit(1)

431

432 wakeup_time_stored = int(wakeup_time_stored_str)

433 try:

434 logging.debug(

435 "Wakeup timestamp: {0} ({1})".format(

436 wakeup_time_stored,

437 datetime.fromtimestamp(wakeup_time_stored).strftime(

438 "%c"

439 ),

440 )

441 )

442 except ValueError as e:

443 logging.error(e)

444 sys.exit(1)

445

446 if (abs(wakeup_time_utc - wakeup_time_stored) > 1) and (

447 abs(wakeup_time_local - wakeup_time_stored) > 1

448 ):

449 logging.error("Wakeup time not stored correctly")

450 sys.exit(1)

451

452 with open(cls.RTC_FILENAME) as rtc_file:

453 separator_regex = re.compile(r"\s+:\s+")

454 rtc_data = dict(

455 [separator_regex.split(line.rstrip()) for line in rtc_file]

456 )

457 logging.debug(

458 "RTC data:\n{0}".format(

459 "\n".join(

460 [

461 "- {0}: {1}".format(*pair)

462 for pair in rtc_data.items()

463 ]

464 )

465 )

466 )

467

468 # Verify wakeup time has been set properly

469 # by looking into the alarm_IRQ and alrm_date field

470 if rtc_data["alarm_IRQ"] != "yes":

471 logging.error(

472 "alarm_IRQ not set properly: {0}".format(

473 rtc_data["alarm_IRQ"]

474 )

475 )

476 sys.exit(1)

477

478 if "*" in rtc_data["alrm_date"]:

479 logging.error(

480 "alrm_date not set properly: {0}".format(

481 rtc_data["alrm_date"]

482 )

483 )

484 sys.exit(1)

485

486

487 class Command:

488 """

489 Simple subprocess.Popen wrapper to run shell commands

490 and log their output

491 """

492

493 def __init__(self, command_str, verbose=True):

494 self.command_str = command_str

495 self.verbose = verbose

496

497 self.process = None

498 self.stdout = None

499 self.stderr = None

500 self.time = None

501

502 def run(self):

503 """

504 Execute shell command and return output and status

505 """

506 logging.debug("Executing: {0!r}...".format(self.command_str))

507

508 self.process = subprocess.Popen(

509 self.command_str,

510 shell=True,

511 stdout=subprocess.PIPE,

512 stderr=subprocess.PIPE,

513 )

514 start = datetime.now()

515 result = self.process.communicate()

516 end = datetime.now()

517 self.time = end - start

518

519 if self.verbose:

520 stdout, stderr = result

521 message = [

522 "Output:\n"

523 "- returncode:\n{0}".format(self.process.returncode)

524 ]

525 if stdout:

526 if type(stdout) is bytes:

527 stdout = stdout.decode("utf-8", "ignore")

528 message.append("- stdout:\n{0}".format(stdout))

529 if stderr:

530 if type(stderr) is bytes:

531 stderr = stderr.decode("utf-8", "ignore")

532 message.append("- stderr:\n{0}".format(stderr))

533 logging.debug("\n".join(message))

534

535 self.stdout = stdout

536 self.stderr = stderr

537

538 return self

539

540

541 class CountdownDialog(Gtk.Dialog):

542 """

543 Dialog that shows the amount of progress in the reboot test

544 and lets the user cancel it if needed

545 """

546

547 def __init__(

548 self,

549 pm_operation,

550 pm_delay,

551 hardware_delay,

552 iterations,

553 iterations_count,

554 ):

555 self.pm_operation = pm_operation

556 title = "{0} test".format(pm_operation.capitalize())

557

558 buttons = (

559 Gtk.STOCK_CANCEL,

560 Gtk.ResponseType.CANCEL,

561 )

562 super(CountdownDialog, self).__init__(title=title, buttons=buttons)

563 self.set_default_response(Gtk.ResponseType.CANCEL)

564 self.set_resizable(False)

565 self.set_position(Gtk.WindowPosition.CENTER)

566

567 progress_bar = Gtk.ProgressBar()

568 progress_bar.set_fraction(iterations / float(iterations_count))

569 progress_bar.set_text("{0}/{1}".format(iterations, iterations_count))

570 progress_bar.set_show_text(True)

571 self.vbox.pack_start(progress_bar, True, True, 0)

572

573 operation_event = {

574 "template": (

575 "Next {0} in {{time}} seconds...".format(self.pm_operation)

576 ),

577 "timeout": pm_delay,

578 }

579 hardware_info_event = {

580 "template": "Gathering hardware information in {time} seconds...",

581 "timeout": hardware_delay,

582 "callback": self.on_hardware_info_timeout_cb,

583 }

584 system_info_event = {

585 "template": "Gathering system information in {time} seconds...",

586 "timeout": 2,

587 "callback": self.on_system_info_timeout_cb,

588 }

589

590 if iterations == 0:

591 # In first iteration, gather hardware and system information

592 # directly and perform pm-operation

593 self.on_hardware_info_timeout_cb()

594 self.on_system_info_timeout_cb()

595 self.events = [operation_event]

596 elif iterations < iterations_count:

597 # In last iteration, wait before gathering hardware information

598 # and perform pm-operation

599 self.events = [

600 operation_event,

601 hardware_info_event,

602 system_info_event,

603 ]

604 else:

605 # In last iteration, wait before gathering hardware information

606 # and finish the test

607 self.events = [hardware_info_event, system_info_event]

608

609 self.label = Gtk.Label()

610 self.vbox.pack_start(self.label, True, True, 0)

611 self.show_all()

612

613 def run(self):

614 """

615 Set label text and run dialog

616 """

617 self.schedule_next_event()

618 response = super(CountdownDialog, self).run()

619 self.destroy()

620 self.show()

621

622 if response != Gtk.ResponseType.ACCEPT:

623 raise TestCancelled()

624

625 def schedule_next_event(self):

626 """

627 Schedule next timed event

628 """

629 if self.events:

630 self.event = self.events.pop()

631 self.timeout_counter = self.event.get("timeout", 0)

632 self.label.set_text(

633 self.event["template"].format(time=self.timeout_counter)

634 )

635 GObject.timeout_add_seconds(1, self.on_timeout_cb)

636 else:

637 # Return Accept response

638 # if there are no other events scheduled

639 self.response(Gtk.ResponseType.ACCEPT)

640

641 def on_timeout_cb(self):

642 """

643 Set label properly and use callback method if needed

644 """

645 if self.timeout_counter > 0:

646 self.label.set_text(

647 self.event["template"].format(time=self.timeout_counter)

648 )

649 self.timeout_counter -= 1

650 return True

651

652 # Call calback if defined

653 callback = self.event.get("callback")

654 if callback:

655 callback()

656

657 # Schedule next event if needed

658 self.schedule_next_event()

659

660 return False

661

662 def on_hardware_info_timeout_cb(self):

663 """

664 Gather hardware information and print it to logs

665 """

666 logging.info("Gathering hardware information...")

667 logging.debug(

668 "Networking:\n"

669 "{network}\n"

670 "{ethernet}\n"

671 "{ifconfig}\n"

672 "{iwconfig}".format(

673 network=(Command("lspci | grep Network").run().stdout),

674 ethernet=(Command("lspci | grep Ethernet").run().stdout),

675 ifconfig=(

676 Command(r"ifconfig -a | grep -A1 '^\w'").run().stdout

677 ),

678 iwconfig=(Command(r"iwconfig | grep -A1 '^\w'").run().stdout),

679 )

680 )

681 logging.debug(

682 "Bluetooth Device:\n"

683 "{hciconfig}".format(

684 hciconfig=(

685 Command(r"hciconfig -a " r"| grep -A2 '^\w'").run().stdout

686 )

687 )

688 )

689 logging.debug(

690 "Video Card:\n"

691 "{lspci}".format(lspci=Command("lspci | grep VGA").run().stdout)

692 )

693 logging.debug(

694 "Touchpad and Keyboard:\n"

695 "{xinput}".format(

696 xinput=Command("xinput list --name-only | sort").run().stdout

697 )

698 )

699 logging.debug(

700 "Pulse Audio Sink:\n"

701 "{pactl_sink}".format(

702 pactl_sink=(Command("pactl list | grep Sink").run().stdout)

703 )

704 )

705 logging.debug(

706 "Pulse Audio Source:\n"

707 "{pactl_source}".format(

708 pactl_source=(Command("pactl list | grep Source").run().stdout)

709 )

710 )

711

712 # Check kernel logs using firmware test suite

713 command = Command("fwts -r stdout klog oops").run()

714 if command.process.returncode != 0:

715 # Don't abort the test loop,

716 # errors can be retrieved by pm_log_check.py

717 logging.error("Problem found in logs by fwts")

718

719 def on_system_info_timeout_cb(self):

720 """

721 Gather system information and print it to logs

722 """

723 logging.info("Gathering system information...")

724 # FIXME: Commented out as it created huge log files

725 # during stress tests.

726 # logging.debug('--- beginning of dmesg ---')

727 # logging.debug(Command('dmesg').run().stdout)

728 # logging.debug('--- end of dmesg ---')

729 # logging.debug('--- beginning of syslog ---')

730 # logging.debug(Command('cat /var/log/syslog').run().stdout)

731 # logging.debug('--- end of syslog ---')

732

733

734 class MessageDialog:

735 """

736 Simple wrapper aroung Gtk.MessageDialog

737 """

738

739 def __init__(self, title, message, type=Gtk.MessageType.INFO):

740 self.title = title

741 self.message = message

742 self.type = type

743

744 def run(self):

745 dialog = Gtk.MessageDialog(

746 buttons=Gtk.ButtonsType.OK,

747 message_format=self.message,

748 type=self.type,

749 )

750 logging.info(self.message)

751 dialog.set_title(self.title)

752 dialog.run()

753 dialog.destroy()

754

755

756 class AutoLoginConfigurator:

757 """

758 Enable/disable autologin configuration

759 to make sure that reboot test will work properly

760 """

761

762 def __init__(self, user=None):

763 self.user = user

764 self.config_filename = "/etc/lightdm/lightdm.conf"

765 self.template = """

766[SeatDefaults]

767 greeter-session=unity-greeter

768 user-session=ubuntu

769 autologin-user={username}

770 autologin-user-timeout=0

771"""

772 if os.path.exists("/etc/gdm3/custom.conf"):

773 self.config_filename = "/etc/gdm3/custom.conf"

774 self.parser = ConfigParser()

775 self.parser.optionxform = str

776 self.parser.read(self.config_filename)

777 self.parser.set("daemon", "AutomaticLoginEnable", "True")

778 if self.user:

779 self.parser.set("daemon", "AutomaticLogin", self.user)

780

781 def enable(self):

782 """

783 Make sure user will autologin in next reboot

784 """

785 logging.debug("Enabling autologin for this user...")

786 if os.path.exists(self.config_filename):

787 for backup_filename in self.generate_backup_filename():

788 if not os.path.exists(backup_filename):

789 shutil.copyfile(self.config_filename, backup_filename)

790 shutil.copystat(self.config_filename, backup_filename)

791 break

792

793 with open(self.config_filename, "w") as f:

794 if self.config_filename == "/etc/lightdm/lightdm.conf":

795 f.write(self.template.format(username=self.user))

796 elif self.config_filename == "/etc/gdm3/custom.conf":

797 self.parser.write(f)

798

799 def disable(self):

800 """

801 Remove latest configuration file

802 and use the same configuration that was in place

803 before running the test

804 """

805 logging.debug("Restoring autologin configuration...")

806 backup_filename = None

807 for filename in self.generate_backup_filename():

808 if not os.path.exists(filename):

809 break

810 backup_filename = filename

811

812 if backup_filename:

813 shutil.copy(backup_filename, self.config_filename)

814 os.remove(backup_filename)

815 else:

816 os.remove(self.config_filename)

817

818 def generate_backup_filename(self):

819 backup_filename = self.config_filename + ".bak"

820 yield backup_filename

821

822 index = 0

823 while True:

824 index += 1

825 backup_filename = self.config_filename + ".bak.{0}".format(index)

826 yield backup_filename

827

828

829 class SudoersConfigurator:

830 """

831 Enable/disable reboot test to be executed as root

832 to make sure that reboot test works properly

833 """

834

835 MARK = "# Automatically added by pm.py"

836 SUDOERS = "/etc/sudoers"

837

838 def __init__(self, user=None):

839 self.user = user

840

841 def enable(self):

842 """

843 Make sure that user will be allowed to execute reboot test as root

844 """

845 logging.debug("Enabling user to execute test as root...")

846 command = (

847 "sed -i -e '$a{mark}\\n"

848 "{user} ALL=NOPASSWD: /usr/bin/python3' "

849 "{filename}".format(

850 mark=self.MARK, user=self.user, filename=self.SUDOERS

851 )

852 )

853

854 Command(command, verbose=False).run()

855

856 def disable(self):

857 """

858 Revert sudoers configuration changes

859 """

860 logging.debug("Restoring sudoers configuration...")

861 command = ("sed -i -e '/{mark}/,+1d' " "{filename}").format(

862 mark=self.MARK, filename=self.SUDOERS

863 )

864 Command(command, verbose=False).run()

865

866

867 class AutoStartFile:

868 """

869 Generate autostart file contents and write it to proper location

870 """

871

872 TEMPLATE = """

873[Desktop Entry]

874 Name={pm_operation} test

875 Comment=Verify {pm_operation} works properly

876 Exec=sudo {script} -r {repetitions} -w {wakeup} --hardware-delay {hardware_delay} --pm-delay {pm_delay} --min-pm-time {min_pm_time} --max-pm-time {max_pm_time} --append --total {total} --start {start} --pm-timestamp {pm_timestamp} {silent} --log-level={log_level} --log-dir={log_dir} --suspends-before-reboot={suspend_cycles} --checkbox-respawn-cmd={checkbox_respawn} {check_hardware} {fwts} {pm_operation}

877 Type=Application

878 X-GNOME-Autostart-enabled=true

879 Hidden=false

880""" # noqa: E501

881

882 def __init__(self, args, user=None):

883 self.args = args

884 self.user = user

885

886 # Generate desktop filename

887 # based on environment variables

888 username = self.user

889 default_config_directory = os.path.expanduser(

890 "~{0}/.config".format(username)

891 )

892 config_directory = os.getenv(

893 "XDG_CONFIG_HOME", default_config_directory

894 )

895 autostart_directory = os.path.join(config_directory, "autostart")

896 if not os.path.exists(autostart_directory):

897 os.makedirs(autostart_directory)

898 user_id = os.getenv("PKEXEC_UID") or os.getenv("SUDO_UID")

899 group_id = os.getenv("PKEXEC_UID") or os.getenv("SUDO_GID")

900 if user_id:

901 os.chown(config_directory, int(user_id), int(group_id))

902 os.chown(autostart_directory, int(user_id), int(group_id))

903

904 basename = "{0}.desktop".format(os.path.basename(__file__))

905 self.desktop_filename = os.path.join(autostart_directory, basename)

906

907 def write(self):

908 """

909 Write autostart file to execute the script on startup

910 """

911 logging.debug(

912 "Writing desktop file ({0!r})...".format(self.desktop_filename)

913 )

914 snap_name = os.getenv("SNAP_NAME")

915 if snap_name:

916 script = "/snap/bin/{}.pm-test".format(snap_name)

917 else:

918 script = "/usr/bin/python3 {}".format(os.path.realpath(__file__))

919 contents = self.TEMPLATE.format(

920 script=script,

921 repetitions=self.args.repetitions - 1,

922 wakeup=self.args.wakeup,

923 hardware_delay=self.args.hardware_delay,

924 pm_delay=self.args.pm_delay,

925 min_pm_time=self.args.min_pm_time,

926 max_pm_time=self.args.max_pm_time,

927 total=self.args.total,

928 start=self.args.start,

929 pm_timestamp=int(time()),

930 silent="--silent" if self.args.silent else "",

931 log_level=self.args.log_level_str,

932 log_dir=self.args.log_dir,

933 fwts="--fwts" if self.args.fwts else "",

934 suspend_cycles=self.args.suspends_before_reboot,

935 pm_operation=self.args.pm_operation,

936 checkbox_respawn=self.args.checkbox_respawn_cmd,

937 check_hardware=(

938 "--check-hardware-list"

939 if self.args.check_hardware_list

940 else ""

941 ),

942 )

943 logging.debug(contents)

944

945 with open(self.desktop_filename, "w") as f:

946 f.write(contents)

947

948 def remove(self):

949 """

950 Remove autostart file to avoid executing the script on startup

951 """

952 if os.path.exists(self.desktop_filename):

953 logging.debug(

954 "Removing desktop file ({0!r})...".format(

955 self.desktop_filename

956 )

957 )

958 os.remove(self.desktop_filename)

959

960

961 class LoggingConfiguration:

962 @classmethod

963 def set(cls, log_level, log_filename, append):

964 """

965 Configure a rotating file logger

966 """

967 logger = logging.getLogger()

968 logger.setLevel(logging.DEBUG)

969

970 # Log to sys.stdout using log level passed through command line

971 if log_level != logging.NOTSET:

972 log_handler = logging.StreamHandler(sys.stdout)

973 formatter = logging.Formatter("%(levelname)-8s %(message)s")

974 log_handler.setFormatter(formatter)

975 log_handler.setLevel(log_level)

976 logger.addHandler(log_handler)

977

978 # Log to rotating file using DEBUG log level

979 log_handler = logging.handlers.RotatingFileHandler(

980 log_filename, mode="a+", backupCount=3

981 )

982 formatter = logging.Formatter(

983 "%(asctime)s %(levelname)-8s " "%(message)s"

984 )

985 log_handler.setFormatter(formatter)

986 log_handler.setLevel(logging.DEBUG)

987 logger.addHandler(log_handler)

988

989 if not append:

990 # Create a new log file on every new

991 # (i.e. not scheduled) invocation

992 log_handler.doRollover()

993

994

995 class MyArgumentParser:

996 """

997 Command-line argument parser

998 """

999

1000 def __init__(self):

1001 """

1002 Create parser object

1003 """

1004 pm_operations = ("poweroff", "reboot")

1005 description = "Run power management operation as many times as needed"

1006 epilog = (

1007 "Unknown arguments will be passed "

1008 "to the underlying command: poweroff or reboot."

1009 )

1010 parser = ArgumentParser(description=description, epilog=epilog)

1011 parser.add_argument(

1012 "-r",

1013 "--repetitions",

1014 type=int,

1015 default=1,

1016 help=(

1017 "Number of times that the power management "

1018 "operation has to be repeated "

1019 "(%(default)s by default)"

1020 ),

1021 )

1022 parser.add_argument(

1023 "-w",

1024 "--wakeup",

1025 type=int,

1026 default=60,

1027 help=(

1028 "Timeout in seconds for the wakeup alarm "

1029 "(%(default)s by default). "

1030 "Note: wakeup alarm won't be scheduled "

1031 "for reboot."

1032 ),

1033 )

1034 parser.add_argument(

1035 "--min-pm-time",

1036 dest="min_pm_time",

1037 type=int,

1038 default=0,

1039 help=(

1040 "Minimum time in seconds that "

1041 "it should take the power management "

1042 "operation each cycle (0 for reboot and "

1043 "wakeup time minus two seconds "

1044 "for the other power management operations "

1045 "by default)"

1046 ),

1047 )

1048 parser.add_argument(

1049 "--max-pm-time",

1050 dest="max_pm_time",

1051 type=int,

1052 default=300,

1053 help=(

1054 "Maximum time in seconds "

1055 "that it should take "

1056 "the power management operation each cycle "

1057 "(%(default)s by default)"

1058 ),

1059 )

1060 parser.add_argument(

1061 "--pm-delay",

1062 dest="pm_delay",

1063 type=int,

1064 default=5,

1065 help=(

1066 "Delay in seconds "

1067 "after hardware information "

1068 "has been gathered and before executing "

1069 "the power management operation "

1070 "(%(default)s by default)"

1071 ),

1072 )

1073 parser.add_argument(

1074 "--hardware-delay",

1075 dest="hardware_delay",

1076 type=int,

1077 default=30,

1078 help=(

1079 "Delay in seconds before gathering hardware "

1080 "information (%(default)s by default)"

1081 ),

1082 )

1083 parser.add_argument(

1084 "--silent",

1085 action="store_true",

1086 help=(

1087 "Don't display any dialog "

1088 "when test is complete "

1089 "to let the script be used "

1090 "in automated tests"

1091 ),

1092 )

1093 log_levels = [

1094 "notset",

1095 "debug",

1096 "info",

1097 "warning",

1098 "error",

1099 "critical",

1100 ]

1101 parser.add_argument(

1102 "--log-level",

1103 dest="log_level_str",

1104 default="info",

1105 choices=log_levels,

1106 help=(

1107 "Log level. "

1108 "One of {0} or {1} (%(default)s by default)".format(

1109 ", ".join(log_levels[:-1]), log_levels[-1]

1110 )

1111 ),

1112 )

1113 parser.add_argument(

1114 "--log-dir",

1115 dest="log_dir",

1116 default="/var/log",

1117 help=("Path to the directory to store log files"),

1118 )

1119 parser.add_argument(

1120 "pm_operation",

1121 choices=pm_operations,

1122 help=(

1123 "Power management operation to be performed "

1124 "(one of {0} or {1!r})".format(

1125 ", ".join(map(repr, pm_operations[:-1])), pm_operations[-1]

1126 )

1127 ),

1128 )

1129

1130 # Test timestamps

1131 parser.add_argument("--start", type=int, default=0, help=SUPPRESS)

1132 parser.add_argument(

1133 "--pm-timestamp",

1134 dest="pm_timestamp",

1135 type=int,

1136 default=0,

1137 help=SUPPRESS,

1138 )

1139

1140 # Append to log on subsequent startups

1141 parser.add_argument(

1142 "--append", action="store_true", default=False, help=SUPPRESS

1143 )

1144

1145 # Total number of iterations initially passed through the command line

1146 parser.add_argument("--total", type=int, default=0, help=SUPPRESS)

1147

1148 # suspend cycles before reboot

1149 parser.add_argument(

1150 "--suspends-before-reboot",

1151 type=int,

1152 default=0,

1153 help=(

1154 "How many cycles of suspend/resume to run"

1155 "before each reboot or poweroff"

1156 ),

1157 )

1158

1159 # use fwts for suspend tests

1160 parser.add_argument(

1161 "--fwts",

1162 action="store_true",

1163 help=("Use fwts " "when doing the suspend tests"),

1164 )

1165 parser.add_argument(

1166 "--checkbox-respawn-cmd",

1167 type=str,

1168 help=(

1169 "path to a file telling how to return to checkbox after the"

1170 " test is done"

1171 ),

1172 default="",

1173 )

1174 parser.add_argument(

1175 "--check-hardware-list",

1176 action="store_true",

1177 help=(

1178 "Look for changes in the list of devices "

1179 "after each PM action"

1180 ),

1181 default=False,

1182 )

1183 self.parser = parser

1184

1185 def parse(self):

1186 """

1187 Parse command-line arguments

1188 """

1189 args, extra_args = self.parser.parse_known_args()

1190 args.log_level = getattr(logging, args.log_level_str.upper())

1191

1192 # Total number of repetitions

1193 # is the number of repetitions passed through the command line

1194 # the first time the script is executed

1195 if not args.total:

1196 args.total = args.repetitions

1197

1198 # Test start time automatically set on first iteration

1199 if not args.start:

1200 args.start = int(time())

1201

1202 # Wakeup time set to 0 for 'reboot'

1203 # since wakeup alarm won't be scheduled

1204 if args.pm_operation == "reboot":

1205 args.wakeup = 0

1206 args.min_pm_time = 0

1207

1208 # Minimum time for each power management operation

1209 # is set to the wakeup time

1210 if not args.min_pm_time:

1211 min_pm_time = args.wakeup - 2

1212 if min_pm_time < 0:

1213 min_pm_time = 0

1214 args.min_pm_time = min_pm_time

1215

1216 # Log filename shows clearly the type of test (pm_operation)

1217 # and the times it was repeated (repetitions)

1218 args.log_filename = os.path.join(

1219 args.log_dir,

1220 "{0}.{1}.{2}.log".format(

1221 os.path.splitext(os.path.basename(__file__))[0],

1222 args.pm_operation,

1223 args.total,

1224 ),

1225 )

1226 return args, extra_args

1227

1228

1229 if __name__ == "__main__":

1230 sys.exit(main())