usb_test.py

import os, re, pathlib, subprocess, csv

if os.getuid() != 0:
    print("You need to be root to run this script")
    exit(0)

mounted = False

with open("/proc/mounts", "r") as f:
    for line in f.readlines():
        if line.split()[1] == "/media/usbstick":
            mounted = True

if mounted:
    print("USB device found")
else:
    print("USB device not found")
    devs = enumerate(pathlib.Path("/dev/").glob("sd*"))
    names = []
    for i, dev in devs:
        print(str(i) + ". " + dev.name)
        names.append(str(dev))
    print("Choose USB device")
    usb_device = input()
    if not os.path.exists("/media/usbstick/"):
        os.mkdir("/media/usbstick/")
    os.system("mount " + names[int(usb_device)] + " /media/usbstick/")
    print("USB device mounted")

seq = input("Read or write? [r/w]").lower()
if not seq in ["r", "w"]:
    print("Invalid sequence")
    exit(0)

results = []

for i in range(1, 30):
    out = subprocess.check_output(
        [
            "sh",
            "-c",
            f"dd {'if' if seq == "w" else 'of'}={'/dev/zero' if seq == "w" else '/dev/null'} {'of' if seq == "w" else 'if'}=/media/usbstick/test bs=1M count=512 oflag=dsync {'oflag' if seq == "w" else 'iflag'}=direct",
        ],
        stderr=subprocess.STDOUT,
    )
    res = re.search("[\\d.]* s", out.decode("utf-8"))
    print(out.decode("utf-8"))
    if res:
        results.append(res.group(0))

clean_results = []

for res in results:
    n_res = ""
    for char in res:
        if char.isdigit() or char == ".":
            n_res += char
    clean_results.append((536870912 / float(n_res)) / 1000000)

avg = sum(clean_results) / len(clean_results)
max = max(clean_results)
min = min(clean_results)
diff = max - min

print("Average speed: " + str(avg) + " MB/s")
print("Maximum speed: " + str(max) + " MB/s")
print("Minimum speed: " + str(min) + " MB/s")
print("Difference: " + str(diff) + " MB/s")

with open("results.csv", "w") as f:
    writer = csv.writer(f)
    writer.writerow(["Speed (MB/s)", "Average", "Maximum", "Minimum", "Difference"])
    for i, res in enumerate(clean_results):
        if i == 0:
            writer.writerow([res, avg, max, min, diff])
            continue

        writer.writerow([res])

if mounted:
    os.system("umount /media/usbstick/")