1#!/usr/bin/env python3

2

3 import argparse

4 import collections

5 import dbus

6 import hashlib

7 import logging

8 import os

9 import platform

10 import re

11 import shlex

12 import subprocess

13 import sys

14 import tempfile

15 import time

16

17 import gi

18

19 gi.require_version("GUdev", "1.0")

20 from gi.repository import GUdev # noqa: E402

21

22 from checkbox_support.dbus import connect_to_system_bus # noqa: E402

23 from checkbox_support.dbus.udisks2 import ( # noqa: E402

24 UDISKS2_BLOCK_INTERFACE,

25 UDISKS2_DRIVE_INTERFACE,

26 UDISKS2_FILESYSTEM_INTERFACE,

27 UDISKS2_LOOP_INTERFACE,

28 UDisks2Model,

29 UDisks2Observer,

30 is_udisks2_supported,

31 lookup_udev_device,

32 map_udisks1_connection_bus,

33)

34 from checkbox_support.heuristics.udisks2 import is_memory_card # noqa: E402

35 from checkbox_support.helpers.human_readable_bytes import ( # noqa: E402

36 HumanReadableBytes,

37)

38 from checkbox_support.parsers.udevadm import ( # noqa: E402

39 CARD_READER_RE,

40 GENERIC_RE,

41 FLASH_RE,

42 find_pkname_is_root_mountpoint,

43) # noqa: E402

44 from checkbox_support.udev import get_interconnect_speed # noqa: E402

45 from checkbox_support.udev import get_udev_block_devices # noqa: E402

46 from checkbox_support.udev import get_udev_xhci_devices # noqa: E402

47

48

49 class ActionTimer:

50 """Class to implement a simple timer"""

51

52 def __enter__(self):

53 self.start = time.time()

54 return self

55

56 def __exit__(self, *args):

57 self.stop = time.time()

58 self.interval = self.stop - self.start

59

60

61 class RandomData:

62 """Class to create data files"""

63

64 def __init__(self, size):

65 self.tfile = tempfile.NamedTemporaryFile(delete=False)

66 self.path = ""

67 self.name = ""

68 self.path, self.name = os.path.split(self.tfile.name)

69 self._write_test_data_file(size)

70

71 def _generate_test_data(self):

72 seed = "104872948765827105728492766217823438120"

73 phrase = """

74 Lorem ipsum dolor sit amet, consectetuer adipiscing elit, sed diam

75 nonummy nibh euismod tincidunt ut laoreet dolore magna aliquam erat

76 volutpat. Ut wisi enim ad minim veniam, quis nostrud exerci tation

77 ullamcorper suscipit lobortis nisl ut aliquip ex ea commodo consequat.

78 Duis autem vel eum iriure dolor in hendrerit in vulputate velit esse

79 molestie consequat, vel illum dolore eu feugiat nulla facilisis at vero

80 eros et accumsan et iusto odio dignissim qui blandit praesent luptatum

81 zzril delenit augue duis dolore te feugait nulla facilisi.

82 """

83 words = phrase.replace("\n", "").split()

84 word_deque = collections.deque(words)

85 seed_deque = collections.deque(seed)

86 while True:

87 yield " ".join(list(word_deque))

88 word_deque.rotate(int(seed_deque[0]))

89 seed_deque.rotate(1)

90

91 def _write_test_data_file(self, size):

92 data = self._generate_test_data()

93 while os.path.getsize(self.tfile.name) < size:

94 self.tfile.write(next(data).encode("UTF-8"))

95 return self

96

97

98 def md5_hash_file(path):

99 md5 = hashlib.md5()

100 try:

101 with open(path, "rb") as stream:

102 while True:

103 data = stream.read(8192)

104 if not data:

105 break

106 md5.update(data)

107 except IOError as exc:

108 logging.error("unable to checksum %s: %s", path, exc)

109 return None

110 else:

111 return md5.hexdigest()

112

113

114 def on_ubuntucore():

115 """

116 Check if running from on ubuntu core

117 """

118 snap = os.getenv("SNAP")

119 if snap:

120 with open(os.path.join(snap, "meta/snap.yaml")) as f:

121 for line in f.readlines():

122 if line == "confinement: classic\n":

123 return False

124 return True

125 return False

126

127

128 class DiskTest:

129 """Class to contain various methods for testing removable disks"""

130

131 def __init__(self, device, memorycard, lsblkcommand):

132 self.rem_disks = {} # mounted before the script running

133 self.rem_disks_nm = {} # not mounted before the script running

134 self.rem_disks_memory_cards = {}

135 self.rem_disks_memory_cards_nm = {}

136 self.rem_disks_speed = {}

137 # LP: #1313581, TODO: extend to be rem_disks_driver

138 self.rem_disks_xhci = {}

139 self.data = ""

140 self.lsblk = ""

141 self.device = device

142 self.memorycard = memorycard

143 self._run_lsblk(lsblkcommand)

144 self._probe_disks()

145

146 def read_file(self, source):

147 with open(source, "rb") as infile:

148 try:

149 self.data = infile.read()

150 except IOError as exc:

151 logging.error("Unable to read data from %s: %s", source, exc)

152 return False

153 else:

154 return True

155

156 def write_file(self, data, dest):

157 try:

158 outfile = open(dest, "wb", 0)

159 except OSError as exc:

160 logging.error("Unable to open %s for writing.", dest)

161 logging.error(" %s", exc)

162 return False

163 with outfile:

164 try:

165 outfile.write(self.data)

166 except IOError as exc:

167 logging.error("Unable to write data to %s: %s", dest, exc)

168 return False

169 else:

170 outfile.flush()

171 os.fsync(outfile.fileno())

172 return True

173

174 def clean_up(self, target):

175 try:

176 os.unlink(target)

177 except OSError as exc:

178 logging.error("Unable to remove tempfile %s", target)

179 logging.error(" %s", exc)

180

181 def _find_parent(self, device):

182 if self.lsblk:

183 pattern = re.compile(

184 r'KNAME="(?P<KNAME>.*)" '

185 r'TYPE="(?P<TYPE>.*)" '

186 r'MOUNTPOINT="(?P<MOUNTPOINT>.*)"'

187 )

188 for line in self.lsblk.splitlines():

189 m = pattern.match(line)

190 if m and device.startswith(m.group("KNAME")):

191 return m.group("KNAME")

192 return False

193

194 def _run_lsblk(self, lsblkcommand):

195 try:

196 self.lsblk = subprocess.check_output(

197 shlex.split(lsblkcommand), universal_newlines=True

198 )

199 except subprocess.CalledProcessError as exc:

200 raise SystemExit(exc)

201

202 def _probe_disks(self):

203 """

204 Internal method used to probe for available disks

205

206 Indirectly sets:

207 self.rem_disks{,_nm,_memory_cards,_memory_cards_nm,_speed}

208 """

209 if on_ubuntucore():

210 self._probe_disks_udisks2_cli()

211 else:

212 bus, loop = connect_to_system_bus()

213 if is_udisks2_supported(bus):

214 self._probe_disks_udisks2(bus)

215 else:

216 self._probe_disks_udisks1(bus)

217

218 def _probe_disks_udisks2_cli(self):

219 # First we will build up a db of udisks info by scraping the output

220 # of the dump command

221 # TODO: remove the snap prefix when the alias becomes available

222 proc = subprocess.Popen(

223 ["udisks2.udisksctl", "dump"], stdout=subprocess.PIPE

224 )

225 udisks_devices = {}

226 current_bd = None

227 current_interface = None

228 while True:

229 line = proc.stdout.readline().decode(sys.stdout.encoding)

230 if line == "":

231 break

232 if line == "\n":

233 current_bd = None

234 current_interface = None

235 if line.startswith("/org/freedesktop/UDisks2/"):

236 path = line.strip()

237 current_bd = os.path.basename(path).rstrip(":")

238 udisks_devices[current_bd] = {}

239 continue

240 if current_bd is None:

241 continue

242 if line.startswith(" org.freedesktop"):

243 current_interface = line.strip().rstrip(":")

244 udisks_devices[current_bd][current_interface] = {}

245 continue

246 if current_interface is None:

247 continue

248 entry = "".join(c for c in line if c not in "\n\t' ")

249 wanted_keys = (

250 "Device:",

251 "Drive:",

252 "MountPoints:",

253 "Vendor:",

254 "ConnectionBus:",

255 "Model:",

256 "Media:",

257 )

258 for key in wanted_keys:

259 if entry.startswith(key):

260 udisks_devices[current_bd][current_interface][key] = entry[

261 len(key) :

262 ]

263

264 # Now use the populated udisks structure to fill out the API used by

265 # other _probe disks functions

266 for device, interfaces in udisks_devices.items():

267 # iterate over udisks objects that have both filesystem and

268 # block device interfaces

269 if (

270 UDISKS2_FILESYSTEM_INTERFACE in interfaces

271 and UDISKS2_BLOCK_INTERFACE in interfaces

272 ):

273 # To be an IO candidate there must be a drive object

274 drive = interfaces[UDISKS2_BLOCK_INTERFACE].get("Drive:")

275 if drive is None or drive == "/":

276 continue

277 drive_object = udisks_devices[os.path.basename(drive)]

278

279 # Get the connection bus property from the drive interface of

280 # the drive object. This is required to filter out the devices

281 # we don't want to look at now.

282 connection_bus = drive_object[UDISKS2_DRIVE_INTERFACE][

283 "ConnectionBus:"

284 ]

285 desired_connection_buses = set(

286 [

287 map_udisks1_connection_bus(device)

288 for device in self.device

289 ]

290 )

291 # Skip devices that are attached to undesired connection buses

292 if connection_bus not in desired_connection_buses:

293 continue

294

295 dev_file = interfaces[UDISKS2_BLOCK_INTERFACE].get("Device:")

296

297 parent = self._find_parent(dev_file.replace("/dev/", ""))

298 if parent and find_pkname_is_root_mountpoint(

299 parent, self.lsblk

300 ):

301 continue

302

303 # XXX: we actually only scrape the first one currently

304 mount_point = interfaces[UDISKS2_FILESYSTEM_INTERFACE].get(

305 "MountPoints:"

306 )

307 if mount_point == "":

308 mount_point = None

309

310 # We need to skip-non memory cards if we look for memory cards

311 # and vice-versa so let's inspect the drive and use heuristics

312 # to detect memory cards (a memory card reader actually) now.

313 if self.memorycard != is_memory_card(

314 drive_object[UDISKS2_DRIVE_INTERFACE]["Vendor:"],

315 drive_object[UDISKS2_DRIVE_INTERFACE]["Model:"],

316 drive_object[UDISKS2_DRIVE_INTERFACE]["Media:"],

317 ):

318 continue

319

320 if mount_point is None:

321 self.rem_disks_memory_cards_nm[dev_file] = None

322 self.rem_disks_nm[dev_file] = None

323 else:

324 self.rem_disks_memory_cards[dev_file] = mount_point

325 self.rem_disks[dev_file] = mount_point

326

327 # Get the speed of the interconnect that is associated with the

328 # block device we're looking at. This is purely informational

329 # but it is a part of the required API

330 udev_devices = get_udev_block_devices(GUdev.Client())

331 for udev_device in udev_devices:

332 if udev_device.get_device_file() == dev_file:

333 interconnect_speed = get_interconnect_speed(

334 udev_device

335 )

336 if interconnect_speed:

337 self.rem_disks_speed[dev_file] = (

338 interconnect_speed * 10**6

339 )

340 else:

341 self.rem_disks_speed[dev_file] = None

342

343 def _probe_disks_udisks2(self, bus):

344 """

345 Internal method used to probe / discover available disks using udisks2

346 dbus interface using the provided dbus bus (presumably the system bus)

347 """

348 # We'll need udisks2 and udev to get the data we need

349 udisks2_observer = UDisks2Observer()

350 udisks2_model = UDisks2Model(udisks2_observer)

351 udisks2_observer.connect_to_bus(bus)

352 udev_client = GUdev.Client()

353 # Get a collection of all udev devices corresponding to block devices

354 udev_devices = get_udev_block_devices(udev_client)

355 # Get a collection of all udisks2 objects

356 udisks2_objects = udisks2_model.managed_objects

357 # Let's get a helper to simplify the loop below

358

359 def iter_filesystems_on_block_devices():

360 """

361 Generate a collection of UDisks2 object paths that

362 have both the filesystem and block device interfaces

363 """

364 for udisks2_object_path, interfaces in udisks2_objects.items():

365 if (

366 UDISKS2_FILESYSTEM_INTERFACE in interfaces

367 and UDISKS2_BLOCK_INTERFACE in interfaces

368 and UDISKS2_LOOP_INTERFACE not in interfaces

369 ):

370 yield udisks2_object_path

371

372 # We need to know about all IO candidates,

373 # let's iterate over all the block devices reported by udisks2

374 for udisks2_object_path in iter_filesystems_on_block_devices():

375 # Get interfaces implemented by this object

376 udisks2_object = udisks2_objects[udisks2_object_path]

377 # Find the path of the udisks2 object that represents the drive

378 # this object is a part of

379 drive_object_path = udisks2_object[UDISKS2_BLOCK_INTERFACE][

380 "Drive"

381 ]

382 # Lookup the drive object, if any. This can fail when

383 try:

384 drive_object = udisks2_objects[drive_object_path]

385 except KeyError:

386 logging.error(

387 "Unable to locate drive associated with %s",

388 udisks2_object_path,

389 )

390 continue

391 else:

392 drive_props = drive_object[UDISKS2_DRIVE_INTERFACE]

393 # Get the connection bus property from the drive interface of the

394 # drive object. This is required to filter out the devices we don't

395 # want to look at now.

396 connection_bus = drive_props["ConnectionBus"]

397 desired_connection_buses = set(

398 [map_udisks1_connection_bus(device) for device in self.device]

399 )

400 # Skip devices that are attached to undesired connection buses

401 if connection_bus not in desired_connection_buses:

402 continue

403 # Lookup the udev object that corresponds to this object

404 try:

405 udev_device = lookup_udev_device(udisks2_object, udev_devices)

406 except LookupError:

407 logging.error(

408 "Unable to locate udev object that corresponds to: %s",

409 udisks2_object_path,

410 )

411 continue

412 # Get the block device pathname,

413 # to avoid the confusion, this is something like /dev/sdbX

414 dev_file = udev_device.get_device_file()

415 parent = self._find_parent(dev_file.replace("/dev/", ""))

416 if parent and find_pkname_is_root_mountpoint(parent, self.lsblk):

417 continue

418 # Get the list of mount points of this block device

419 mount_points = udisks2_object[UDISKS2_FILESYSTEM_INTERFACE][

420 "MountPoints"

421 ]

422 # Get the speed of the interconnect that is associated with the

423 # block device we're looking at. This is purely informational but

424 # it is a part of the required API

425 interconnect_speed = get_interconnect_speed(udev_device)

426 if interconnect_speed:

427 self.rem_disks_speed[dev_file] = interconnect_speed * 10**6

428 else:

429 self.rem_disks_speed[dev_file] = None

430 # Ensure it is a media card reader if this was explicitly requested

431 drive_is_reader = is_memory_card(

432 drive_props["Vendor"],

433 drive_props["Model"],

434 drive_props["Media"],

435 )

436 if self.memorycard and not drive_is_reader:

437 continue

438 # The if/else test below simply distributes the mount_point to the

439 # appropriate variable, to keep the API requirements. It is

440 # confusing as _memory_cards is variable is somewhat dummy.

441 if mount_points:

442 # XXX: Arbitrarily pick the first of the mount points

443 mount_point = mount_points[0]

444 self.rem_disks_memory_cards[dev_file] = mount_point

445 self.rem_disks[dev_file] = mount_point

446 else:

447 self.rem_disks_memory_cards_nm[dev_file] = None

448 self.rem_disks_nm[dev_file] = None

449

450 def _probe_disks_udisks1(self, bus):

451 """

452 Internal method used to probe / discover available disks using udisks1

453 dbus interface using the provided dbus bus (presumably the system bus)

454 """

455 ud_manager_obj = bus.get_object(

456 "org.freedesktop.UDisks", "/org/freedesktop/UDisks"

457 )

458 ud_manager = dbus.Interface(ud_manager_obj, "org.freedesktop.UDisks")

459 for dev in ud_manager.EnumerateDevices():

460 device_obj = bus.get_object("org.freedesktop.UDisks", dev)

461 device_props = dbus.Interface(device_obj, dbus.PROPERTIES_IFACE)

462 udisks = "org.freedesktop.UDisks.Device"

463 if not device_props.Get(udisks, "DeviceIsDrive"):

464 dev_bus = device_props.Get(udisks, "DriveConnectionInterface")

465 if dev_bus in self.device:

466 parent_model = parent_vendor = ""

467 if device_props.Get(udisks, "DeviceIsPartition"):

468 parent_obj = bus.get_object(

469 "org.freedesktop.UDisks",

470 device_props.Get(udisks, "PartitionSlave"),

471 )

472 parent_props = dbus.Interface(

473 parent_obj, dbus.PROPERTIES_IFACE

474 )

475 parent_model = parent_props.Get(udisks, "DriveModel")

476 parent_vendor = parent_props.Get(udisks, "DriveVendor")

477 parent_media = parent_props.Get(udisks, "DriveMedia")

478 if self.memorycard:

479 if (

480 dev_bus != "sdio"

481 and not FLASH_RE.search(parent_media)

482 and not CARD_READER_RE.search(parent_model)

483 and not GENERIC_RE.search(parent_vendor)

484 ):

485 continue

486 else:

487 if (

488 FLASH_RE.search(parent_media)

489 or CARD_READER_RE.search(parent_model)

490 or GENERIC_RE.search(parent_vendor)

491 ):

492 continue

493 dev_file = str(device_props.Get(udisks, "DeviceFile"))

494 dev_speed = str(

495 device_props.Get(udisks, "DriveConnectionSpeed")

496 )

497 self.rem_disks_speed[dev_file] = dev_speed

498 if len(device_props.Get(udisks, "DeviceMountPaths")) > 0:

499 devPath = str(

500 device_props.Get(udisks, "DeviceMountPaths")[0]

501 )

502 self.rem_disks[dev_file] = devPath

503 self.rem_disks_memory_cards[dev_file] = devPath

504 else:

505 self.rem_disks_nm[dev_file] = None

506 self.rem_disks_memory_cards_nm[dev_file] = None

507

508 def get_disks_xhci(self):

509 """

510 Compare

511 1. the pci slot name of the devices using xhci

512 2. the pci slot name of the disks,

513 which is usb3 disks in this case so far,

514 to make sure the usb3 disk does be on the controller using xhci

515 """

516 # LP: #1378724

517 udev_client = GUdev.Client()

518 # Get a collection of all udev devices corresponding to block devices

519 udev_devices = get_udev_block_devices(udev_client)

520 # Get a collection of all udev devices corresponding to xhci devices

521 udev_devices_xhci = get_udev_xhci_devices(udev_client)

522 if platform.machine() in ("aarch64", "armv7l"):

523 enumerator = GUdev.Enumerator(client=udev_client)

524 udev_devices_xhci = [

525 device

526 for device in enumerator.execute()

527 if (

528 device.get_driver() == "xhci-hcd"

529 or device.get_driver() == "xhci_hcd"

530 )

531 ]

532 for udev_device_xhci in udev_devices_xhci:

533 pci_slot_name = udev_device_xhci.get_property("PCI_SLOT_NAME")

534 xhci_devpath = udev_device_xhci.get_property("DEVPATH")

535 for udev_device in udev_devices:

536 devpath = udev_device.get_property("DEVPATH")

537 if self._compare_pci_slot_from_devpath(devpath, pci_slot_name):

538 self.rem_disks_xhci[

539 udev_device.get_property("DEVNAME")

540 ] = "xhci"

541 if platform.machine() in ("aarch64", "armv7l"):

542 if xhci_devpath in devpath:

543 self.rem_disks_xhci[

544 udev_device.get_property("DEVNAME")

545 ] = "xhci"

546 return self.rem_disks_xhci

547

548 def mount(self):

549 passed_mount = {}

550

551 for key in self.rem_disks_nm:

552 temp_dir = tempfile.mkdtemp()

553 if self._mount(key, temp_dir) != 0:

554 logging.error("can't mount %s", key)

555 else:

556 passed_mount[key] = temp_dir

557

558 if len(self.rem_disks_nm) == len(passed_mount):

559 self.rem_disks_nm = passed_mount

560 return 0

561 else:

562 count = len(self.rem_disks_nm) - len(passed_mount)

563 self.rem_disks_nm = passed_mount

564 return count

565

566 def _mount(self, dev_file, mount_point):

567 return subprocess.call(["mount", dev_file, mount_point])

568

569 def umount(self):

570 errors = 0

571 for disk in self.rem_disks_nm:

572 if not self.rem_disks_nm[disk]:

573 continue

574 if self._umount(disk) != 0:

575 errors += 1

576 logging.error(

577 "can't umount %s on %s", disk, self.rem_disks_nm[disk]

578 )

579 return errors

580

581 def _umount(self, mount_point):

582 # '-l': lazy umount, dealing problem of unable to umount the device.

583 return subprocess.call(["umount", "-l", mount_point])

584

585 def clean_tmp_dir(self):

586 for disk in self.rem_disks_nm:

587 if not self.rem_disks_nm[disk]:

588 continue

589 if not os.path.ismount(self.rem_disks_nm[disk]):

590 os.rmdir(self.rem_disks_nm[disk])

591

592 def _compare_pci_slot_from_devpath(self, devpath, pci_slot_name):

593 # LP: #1334991

594 # a smarter parser to get and validate a pci slot name from DEVPATH

595 # then compare this pci slot name to the other

596 dl = devpath.split("/")

597 s = set([x for x in dl if dl.count(x) > 1])

598 if (

599 (pci_slot_name in dl)

600 and (dl.index(pci_slot_name) < dl.index("block"))

601 and (not (pci_slot_name in s))

602 ):

603 # 1. there is such pci_slot_name

604 # 2. sysfs topology looks like

605 # DEVPATH = ....../pci_slot_name/....../block/......

606 # 3. pci_slot_name should be unique in DEVPATH

607 return True

608 else:

609 return False

610

611

612 def main():

613 parser = argparse.ArgumentParser()

614 parser.add_argument(

615 "device",

616 choices=["usb", "firewire", "sdio", "scsi", "ata_serial_esata"],

617 nargs="+",

618 help=(

619 "The type of removable media "

620 "(usb, firewire, sdio, scsi or ata_serial_esata)"

621 "to test."

622 ),

623 )

624 parser.add_argument(

625 "-l",

626 "--list",

627 action="store_true",

628 default=False,

629 help="List the removable devices and mounting status",

630 )

631 parser.add_argument(

632 "-m",

633 "--min-speed",

634 action="store",

635 default=0,

636 type=int,

637 help="Minimum speed a device must support to be "

638 "considered eligible for being tested (bits/s)",

639 )

640 parser.add_argument(

641 "-p",

642 "--pass-speed",

643 action="store",

644 default=0,

645 type=int,

646 help="Minimum average throughput from all eligible"

647 "devices for the test to pass (MB/s)",

648 )

649 parser.add_argument(

650 "-i",

651 "--iterations",

652 action="store",

653 default="1",

654 type=int,

655 help=(

656 "The number of test cycles to run. One cycle is"

657 "comprised of generating --count data files of "

658 "--size bytes and writing them to each device."

659 ),

660 )

661 parser.add_argument(

662 "-c",

663 "--count",

664 action="store",

665 default="1",

666 type=int,

667 help="The number of random data files to generate",

668 )

669 parser.add_argument(

670 "-s",

671 "--size",

672 action="store",

673 type=HumanReadableBytes,

674 default="1MiB",

675 help=(

676 "The size of the test data file to use. "

677 "You may use SI or IEC suffixes like: 'K', 'M',"

678 "'G', 'T', 'Ki', 'Mi', 'Gi', 'Ti', etc. Default"

679 " is %(default)s"

680 ),

681 )

682 parser.add_argument(

683 "--auto-reduce-size",

684 action="store_true",

685 default=False,

686 help=(

687 "Automatically reduce size to fit in the target"

688 "filesystem. Reducing until fits in 1MiB"

689 ),

690 )

691 parser.add_argument(

692 "-n",

693 "--skip-not-mount",

694 action="store_true",

695 default=False,

696 help=(

697 "skip the removable devices "

698 "which haven't been mounted before the test."

699 ),

700 )

701 parser.add_argument(

702 "--memorycard",

703 action="store_true",

704 help=(

705 "Memory cards devices on bus other than sdio "

706 "require this parameter to identify "

707 "them as such"

708 ),

709 )

710 parser.add_argument(

711 "--driver",

712 choices=["xhci_hcd"],

713 help=(

714 "Detect the driver of the host controller."

715 "Only xhci_hcd for usb3 is supported so far."

716 ),

717 )

718 parser.add_argument(

719 "--lsblkcommand",

720 action="store",

721 type=str,

722 default="lsblk -i -n -P -e 7 -o KNAME,TYPE,MOUNTPOINT",

723 help=(

724 "Command to execute to get lsblk information. "

725 "Only change it if you know what you're doing."

726 ),

727 )

728

729 args = parser.parse_args()

730

731 test = DiskTest(args.device, args.memorycard, args.lsblkcommand)

732

733 # LP:1876966

734 if os.getuid() != 0:

735 print("ERROR: This script must be run as root!")

736 return 1

737

738 errors = 0

739 # If we do have removable drives attached and mounted

740 if len(test.rem_disks) > 0 or len(test.rem_disks_nm) > 0:

741 if args.list: # Simply output a list of drives detected

742 print("-" * 20)

743 print("Removable devices currently mounted:")

744 if args.memorycard:

745 if len(test.rem_disks_memory_cards) > 0:

746 for disk, mnt_point in test.rem_disks_memory_cards.items():

747 print("%s : %s" % (disk, mnt_point))

748 else:

749 print("None")

750

751 print("Removable devices currently not mounted:")

752 if len(test.rem_disks_memory_cards_nm) > 0:

753 for disk in test.rem_disks_memory_cards_nm:

754 print(disk)

755 else:

756 print("None")

757 else:

758 if len(test.rem_disks) > 0:

759 for disk, mnt_point in test.rem_disks.items():

760 print("%s : %s" % (disk, mnt_point))

761 else:

762 print("None")

763

764 print("Removable devices currently not mounted:")

765 if len(test.rem_disks_nm) > 0:

766 for disk in test.rem_disks_nm:

767 print(disk)

768 else:

769 print("None")

770

771 print("-" * 20)

772

773 return 0

774

775 else: # Create a file, copy to disk and compare hashes

776 if args.skip_not_mount:

777 disks_all = test.rem_disks

778 else:

779 # mount those haven't be mounted yet.

780 errors_mount = test.mount()

781

782 if errors_mount > 0:

783 print(

784 "There're total %d device(s) failed at mounting."

785 % errors_mount

786 )

787 errors += errors_mount

788

789 disks_all = dict(

790 list(test.rem_disks.items())

791 + list(test.rem_disks_nm.items())

792 )

793

794 if len(disks_all) > 0:

795 print(

796 "Found the following mounted %s partitions:"

797 % ", ".join(args.device)

798 )

799

800 for disk, mount_point in disks_all.items():

801 supported_speed = test.rem_disks_speed[disk]

802 print(

803 " %s : %s : %s bits/s"

804 % (disk, mount_point, supported_speed),

805 end="",

806 )

807 if args.min_speed and int(args.min_speed) > int(

808 supported_speed

809 ):

810 print(

811 " (Will not test it, speed is below %s bits/s)"

812 % args.min_speed,

813 end="",

814 )

815

816 print("")

817

818 print("-" * 20)

819

820 disks_eligible = {

821 disk: disks_all[disk]

822 for disk in disks_all

823 if not args.min_speed

824 or int(test.rem_disks_speed[disk]) >= int(args.min_speed)

825 }

826 if len(disks_eligible) == 0:

827 logging.error(

828 "No %s disks with speed higher than %s bits/s",

829 args.device,

830 args.min_speed,

831 )

832 return 1

833 write_sizes = []

834 test_files = {}

835 disks_freespace = {}

836 for disk, path in disks_eligible.items():

837 stat = os.statvfs(path)

838 disks_freespace[disk] = stat.f_bfree * stat.f_bsize

839 smallest_freespace = min(disks_freespace.values())

840 smallest_partition = [

841 d

842 for d, v in disks_freespace.items()

843 if v == smallest_freespace

844 ][0]

845 desired_size = args.size

846 if desired_size > smallest_freespace:

847 if args.auto_reduce_size:

848 min_space = HumanReadableBytes("1MiB")

849 if smallest_freespace < min_space:

850 sys.exit(

851 "Not enough space. {} is required on {}".format( # noqa: E501

852 min_space, smallest_partition

853 )

854 )

855 new_size = HumanReadableBytes(

856 int(0.8 * smallest_freespace)

857 )

858 logging.warning(

859 "Automatically reducing test data size"

860 ". {} requested. Reducing to {}.".format(

861 desired_size, new_size

862 )

863 )

864 desired_size = new_size

865 else:

866 sys.exit(

867 "Not enough space. {} is required on {}".format(

868 desired_size, smallest_partition

869 )

870 )

871 # Generate our data file(s)

872 for count in range(args.count):

873 test_files[count] = RandomData(desired_size)

874 write_sizes.append(

875 os.path.getsize(test_files[count].tfile.name)

876 )

877 total_write_size = sum(write_sizes)

878

879 try:

880 # Clear dmesg so we can check for I/O errors later

881 subprocess.check_output(["dmesg", "-C"])

882 for disk, mount_point in disks_eligible.items():

883 print(

884 "%s (Total Data Size / iteration: %0.4f MB):"

885 % (disk, (total_write_size / 1024 / 1024))

886 )

887 iteration_write_size = (

888 (total_write_size * args.iterations) / 1024 / 1024

889 )

890 iteration_write_times = []

891 for iteration in range(args.iterations):

892 target_file_list = []

893 write_times = []

894 for file_index in range(args.count):

895 parent_file = test_files[file_index].tfile.name

896 parent_hash = md5_hash_file(parent_file)

897 target_filename = (

898 test_files[file_index].name

899 + ".%s" % iteration

900 )

901 target_path = mount_point

902 target_file = os.path.join(

903 target_path, target_filename

904 )

905 target_file_list.append(target_file)

906 test.read_file(parent_file)

907 with ActionTimer() as timer:

908 if not test.write_file(

909 test.data, target_file

910 ):

911 logging.error(

912 "Failed to copy %s to %s",

913 parent_file,

914 target_file,

915 )

916 errors += 1

917 continue

918 write_times.append(timer.interval)

919 child_hash = md5_hash_file(target_file)

920 if parent_hash != child_hash:

921 logging.warning(

922 "[Iteration %s] Parent and Child"

923 " copy hashes mismatch on %s!",

924 iteration,

925 target_file,

926 )

927 logging.warning(

928 "\tParent hash: %s", parent_hash

929 )

930 logging.warning(

931 "\tChild hash: %s", child_hash

932 )

933 errors += 1

934 for file in target_file_list:

935 test.clean_up(file)

936 total_write_time = sum(write_times)

937 # avg_write_time = total_write_time / args.count

938 try:

939 avg_write_speed = (

940 (total_write_size / total_write_time)

941 / 1024

942 / 1024

943 )

944 except ZeroDivisionError:

945 avg_write_speed = 0.00

946 finally:

947 iteration_write_times.append(total_write_time)

948 print(

949 "\t[Iteration %s] Average Speed: %0.4f"

950 % (iteration, avg_write_speed)

951 )

952 for iteration in range(args.iterations):

953 iteration_write_time = sum(iteration_write_times)

954 print("\tSummary:")

955 print(

956 "\t\tTotal Data Attempted: %0.4f MB"

957 % iteration_write_size

958 )

959 print(

960 "\t\tTotal Time to write: %0.4f secs"

961 % iteration_write_time

962 )

963 print(

964 "\t\tAverage Write Time: %0.4f secs"

965 % (iteration_write_time / args.iterations)

966 )

967 try:

968 avg_write_speed = (

969 iteration_write_size / iteration_write_time

970 )

971 except ZeroDivisionError:

972 avg_write_speed = 0.00

973 finally:

974 print(

975 "\t\tAverage Write Speed: %0.4f MB/s"

976 % avg_write_speed

977 )

978 finally:

979 for key in range(args.count):

980 test.clean_up(test_files[key].tfile.name)

981 if len(test.rem_disks_nm) > 0:

982 if test.umount() != 0:

983 errors += 1

984 test.clean_tmp_dir()

985 dmesg = subprocess.run(["dmesg"], stdout=subprocess.PIPE)

986 if "I/O error" in dmesg.stdout.decode():

987 logging.error("I/O errors found in dmesg")

988 errors += 1

989

990 if errors > 0:

991 logging.warning(

992 "Completed %s test iterations, but there were"

993 " errors",

994 args.count,

995 )

996 return 1

997 else:

998 # LP: 1313581

999 # Try to figure out whether the disk

1000 # is SuperSpeed USB and using xhci_hcd driver.

1001 if args.driver == "xhci_hcd":

1002 # The speed reported by udisks is sometimes

1003 # less than 5G bits/s, for example,

1004 # it may be 705032705 bits/s

1005 # So using

1006 # 500000000

1007 # = 500 M bits/s

1008 # > 480 M bits/s ( USB 2.0 spec.)

1009 # to make sure that it is higher USB version than 2.0

1010 #

1011 # int() for int(test.rem_disks_speed[disk])

1012 # is necessary

1013 # because the speed value of

1014 # the dictionary rem_disks_speed is

1015 # 1. str or int from _probe_disks_udisks2

1016 # 2. int from _probe_disks_udisks1.

1017 # This is really a mess. : (

1018 print("\t\t--------------------------------")

1019 if 500000000 < int(test.rem_disks_speed[disk]):

1020 print("\t\tDevice Detected: SuperSpeed USB")

1021 # Unlike rem_disks_speed,

1022 # which must has the connect speed

1023 # for each disk devices,

1024 # disk devices may not use xhci as

1025 # controller drivers.

1026 # This will raise KeyError for no

1027 # associated disk device was found.

1028 if test.get_disks_xhci().get(disk, "") != "xhci":

1029 raise SystemExit(

1030 "\t\tDisk does not use xhci_hcd."

1031 )

1032 print("\t\tDriver Detected: xhci_hcd")

1033 else:

1034 # Give it a hint for the detection failure.

1035 # LP: #1362902

1036 print(

1037 (

1038 "\t\tNo SuperSpeed USB using xhci_hcd "

1039 "was detected correctly."

1040 )

1041 )

1042 print(

1043 (

1044 "\t\tHint: please use dmesg to check "

1045 "the system status again."

1046 )

1047 )

1048 return 1

1049 # Pass is not assured

1050 if (

1051 not args.pass_speed

1052 or avg_write_speed >= args.pass_speed

1053 ):

1054 return 0

1055 else:

1056 print(

1057 "FAIL: Average speed was lower than desired "

1058 "pass speed of %s MB/s" % args.pass_speed

1059 )

1060 return 1

1061 else:

1062 logging.error(

1063 "No device being mounted successfully "

1064 "for testing, aborting"

1065 )

1066 return 1

1067

1068 else: # If we don't have removable drives attached and mounted

1069 logging.error("No removable drives were detected, aborting")

1070 return 1

1071

1072

1073 if __name__ == "__main__":

1074 sys.exit(main())