pyqtool

#!/usr/bin/env python3

import re
import os
import sys
import json
import copy
import subprocess
import requests
import resource
from bs4 import BeautifulSoup

PROG_NAME = sys.argv[0]
ORG_DIR = "%s/org/work.org" % os.getenv("HOME")
upstream_keys = ["qemu", "kvm-unit-test", "kvm-unit-tests",
                 "kvm", "linux", "kernel"]
downstream_keys = ["redhat", "rh"]
review_keys = ["review", "reviews", "reviewed"]
misc_keys = ["misc", "miscs"]
done_fmt = re.compile(r'\*+ (DONE|DOING) \[#[ABC]\] (.*)$')
date_fmt = re.compile(r'\* \[.*\] +- +([ABCD]|\d+)$')
key_fmt = re.compile(r'^([^:]+): (.*)$')
bz_fmt = re.compile(r"^\d{7}$")
bz_link_fmt = "https://bugzilla.redhat.com/show_bug.cgi?id="

def shell(cmd):
    proc = subprocess.Popen(cmd, shell=True, stderr=subprocess.PIPE,
                            stdout=subprocess.PIPE)
    out, err = proc.communicate()
    return proc.returncode, out, err

def file_write(path, data):
    try:
        # print "Write data '%s' to file '%s'" % (data, path)
        open(path, "w").write(data)
    except:
        print("Got exception in file_write")

def status_summary_append_entry(summary, entry_in):
    entry = copy.deepcopy(entry_in)
    key = entry["key"]
    if key in downstream_keys:
        summary["downstream"].append(entry)
    elif key in review_keys:
        summary["review"].append(entry["title"])
    elif key in misc_keys:
        summary["misc"].append(entry["title"])
    elif key in upstream_keys and \
         ("PATCH" in entry["title"] or "RFC" in entry["title"]):
        summary["upstream"].append(entry)
    elif entry["subject"].find(bz_link_fmt) == 0:
        entry["subject"] = entry["subject"].replace(bz_link_fmt, "BZ")
        summary["major"].append(entry)
    else:
        summary["major"].append(entry)

def status_summary_dump(summary):
    for entry in summary["upstream"] + summary["downstream"] + summary["major"]:
        print("* %s" % entry["subject"])
        if entry["content"].strip():
            print("\n" + entry["content"].strip("\n") + "\n")
        else:
            print("")
    if summary["review"]:
        print("* Reviewed:")
        print("")
        for line in summary["review"]:
            print("  - %s" % line)
        print("")
    if summary["misc"]:
        print("* Misc:")
        print("")
        for line in summary["misc"]:
            print("  - %s" % line)
        print("")

# input the line, generate a new entry for the task
def status_parse_subject_line(line):
    entry = {}
    line = done_fmt.match(line).groups()[1]
    res = key_fmt.match(line)
    if res:
        key = res.groups()[0].lower()
        title = res.groups()[1]
    else:
        # this is a raw subject, not for patches
        key = ""
        title = line
    entry["subject"] = line
    entry["key"] = key
    entry["title"] = title
    entry["content"] = "\n"     # Empty content
    return entry

def status_summary_append_entry_context(entry, line):
    if line:
        entry["content"] += "  " + line + "\n"
    else:
        # we only keep one newline
        if len(entry["content"]) >= 2 and entry["content"][-2] != "\n\n":
            entry["content"] += "\n"

# return non-null updated entry if parse success, or None which means
# that we stop parsing content here.
def status_parse_content_line(summary, entry, line):
    if not entry or line.startswith("CLOSED:"):
        return entry
    if line == "==END==":
        # stop reading content any more for the reports
        status_summary_append_entry(summary, entry)
        # reset entry
        return {}
    status_summary_append_entry_context(entry, line)
    return entry

score_table = { "A": 100, "B": 75, "C": 50, "D": 25, "F": 0 }

def cmd_generate_status_report(args):
    logfile = open(ORG_DIR)
    entry = {}
    days = 0
    scores = 0
    summary = {
        # Elements: entries
        "upstream": [],
        "downstream": [],
        "major": [],
        # Elements: lines
        "review": [],
        "misc": [],
    }

    while True:
        line = logfile.readline()
        if not line:
            break
        line = line.strip()
        if line.startswith("*"):
            if entry:
                # First, save previous entry if there is
                status_summary_append_entry(summary, entry)
                entry = {}
            date = date_fmt.match(line)
            if date:
                # date line with score
                days += 1
                score = date.groups()[0]
                if score.isdigit():
                    score = int(score)
                else:
                    score = score_table[date.groups()[0]]
                scores += score
            if not done_fmt.match(line):
                # This is possibly a TODO/DOING task
                continue
            # This is a subject line, init entry object
            entry = status_parse_subject_line(line)
        else:
            # This is a content line
            entry = status_parse_content_line(summary, entry, line)

    if entry:
        status_summary_append_entry(summary, entry)

    status_summary_dump(summary)

    print("Scores: %s (%s days)\n" % (scores / days if days else 0, days))

    return 0

PCI_CMD_HELP = """
usage: %s <isolate|stub|recover|status> [bb:dd.ff ...]

Supported PCI subcommands:

isolate: isolate a PCI device using vfio-pci from system
stub:    same as isolate, using pci-stub driver
recover: recover a PCI device by using its generic PCI driver
         (if not specifying any device, will recover all devices)
status:  list current status of PCI devices
""" % PROG_NAME

pci_sysfs_root = "/sys/bus/pci/devices"
pci_tree_cache = None
pci_driver_list = [
    # supported devices with its kernel modules
    {
        "name": "Intel USB xHCI Controller",
        "vendor": "0x8086",
        "device": "0x8c31",
        "module": "xhci_hcd",
    },
    {
        "name": "Intel Ethernet I217-LM",
        "vendor": "0x8086",
        "device": "0x153a",
        "module": "e1000e",
    },
    {
        "name": "Intel Wireless 7260",
        "vendor": "0x8086",
        "device": "0x08b2",
        "module": "iwlwifi",
    },
    {
        "name": "O2 Micro SD Host Controller",
        "vendor": "0x1217",
        "device": "0x8520",
        "module": "sdhci-pci",
    },
    {
        "name": "Intel E1000 Network Device",
        "vendor": "0x8086",
        "device": "0x100e",
        "module": "e1000",
    },
    {
        "name": "Intel 8 Series/C220 Series Chip Audio",
        "vendor": "0x8086",
        "device": "0x8c20",
        "module": "snd_hda_intel",
    }
]

def pci_cmd_help():
    print(PCI_CMD_HELP)
    return 1

def pci_is_bdf(bdf):
    return re.match("^[0-9a-fA-F]{2}:[0-9a-fA-F]{2}.[0-9a-fA-F]$",
                    bdf)

def pci_module_get(vendor, device):
    for driver in pci_driver_list:
        if driver["vendor"] == vendor and driver["device"] == device:
            return driver
    return None

def pci_tree_parse():
    results = []
    dev_list = os.listdir(pci_sysfs_root)
    for dev in dev_list:
        dev_path = pci_sysfs_root + "/" + dev
        vendor = open(dev_path + "/vendor").read().strip()
        device = open(dev_path + "/device").read().strip()
        module = pci_module_get(vendor, device)
        if not module:
            # we don't parse unsupported devices
            continue
        driver_path = dev_path + "/driver"
        if os.access(driver_path, os.F_OK):
            driver = os.readlink(driver_path)
            driver = os.path.basename(driver)
        else:
            driver = ""
        results.append({
            "addr": dev[5:],
            "driver": driver,
            "info": module,
        })
    return results

def pci_tree_get():
    global pci_tree_cache
    if not pci_tree_cache:
        pci_tree_cache = pci_tree_parse()
    return pci_tree_cache

def pci_device_get(addr):
    pci_tree = pci_tree_get()
    for dev in pci_tree:
        if addr == dev["addr"]:
            return dev

def pci_device_list_get_by_driver(whitelist=[], blacklist=[]):
    """Get a list of devices, depending on the specified `whitelist' or
    `blacklist'. Here `whitelist' has higher priority."""
    pci_tree = pci_tree_get()
    result = []
    for dev in pci_tree:
        if dev["driver"] in whitelist:
            # white list first
            result.append(dev)
        elif dev["driver"] in blacklist:
            continue
        else:
            # not in white list, nor in black list, then:
            # 1. if whitelist is empty, we choose it by default
            # 2. if not, we disgard it by default
            if not whitelist:
                result.append(dev)
    return result

def cmd_pci_status(args=[]):
    vfio_list = pci_device_list_get_by_driver(whitelist=["vfio-pci"])
    stub_list = pci_device_list_get_by_driver(whitelist=["pci-stub"])
    other_list = pci_device_list_get_by_driver(blacklist=["vfio-pci",
                                                          "pci-stub"])

    print()
    print("PCI devices that binded to vfio-pci:")
    print()
    for dev in vfio_list:
        print("  %s %s" % (dev["addr"], dev["info"]["name"]))
    if not vfio_list:
        print("  (None)")
    print

    print()
    print("PCI devices that binded to pci-stub:")
    print()
    for dev in stub_list:
        print("  %s %s" % (dev["addr"], dev["info"]["name"]))
    if not stub_list:
        print("  (None)")
    print()

    print("PCI devices that binded to other driver:")
    print()
    for dev in other_list:
        printed = True
        print("  %s %s (%s)" % (dev["addr"], dev["info"]["name"],
                                dev["driver"]))
    if not other_list:
        print("  (None)")
    print()

def pci_device_exists(device):
    dev_tree = pci_tree_get()
    for dev in dev_tree:
        if device == dev["addr"]:
            return True
    return False

def pci_device_list_check(devices):
    # check all devices are valid
    if not devices:
        print("Need to provide device (or device list). Please choose from:")
        cmd_pci_status()
        return 1
    for dev in devices:
        if not pci_is_bdf(dev):
            print("'%s' format wrong (should be BB:DD.FF)" % dev)
            return 1
        if not pci_device_exists(dev):
            print("Device '%s' does not exist" % dev)
            return 1

def pci_driver_bind(driver, addr):
    print("Binding device '%s' to driver '%s'" % (addr, driver))
    dev = pci_device_get(addr)
    if driver in ("vfio-pci", "pci-stub"):
        path_new = "/sys/bus/pci/drivers/%s/new_id" % driver
        data = "%s %s" % (dev["info"]["vendor"][2:],
                        dev["info"]["device"][2:])
        file_write(path_new, data)
    else:
        path_bind = "/sys/bus/pci/drivers/%s/bind" % driver
        file_write(path_bind, "0000:" + addr)

def pci_driver_unbind(driver, addr):
    print("Unbinding device '%s' from driver '%s'" % (addr, driver))
    path = "/sys/bus/pci/drivers/" + driver + "/unbind"
    file_write(path, "0000:" + addr)

def vfio_module_probe():
    shell("modprobe vfio-pci")

def pci_dev_isolate(dev, to_isolate, target="vfio-pci"):
    "Either isolate/recover a PCI device"
    if target not in ("vfio-pci", "pci-stub"):
        print("target '%s' incorrect" % target)
        exit(1)
    if to_isolate:
        # bind it to vfio-pci/pci-stub driver
        if target == "vfio-pci":
            vfio_module_probe()
        if dev["driver"] == target:
            print("Device '%s' already binded to %s, skip" \
                % (dev["addr"], target))
            return
        if dev["driver"]:
            # there is existing driver, unbind
            pci_driver_unbind(dev["driver"], dev["addr"])
        pci_driver_bind(target, dev["addr"])
    else:
        # bind it to generic driver
        module = dev["info"]["module"]
        if dev["driver"] == module:
            print("Device '%s' already binded to driver '%s', skip" \
                % (dev["addr"], module))
            return
        if dev["driver"]:
            pci_driver_unbind(dev["driver"], dev["addr"])
        pci_driver_bind(module, dev["addr"])

def cmd_pci_stub(devices):
    if pci_device_list_check(devices):
        return 1
    dev_tree = pci_tree_get()
    for dev in devices:
        pci_dev_isolate(pci_device_get(dev), True, "pci-stub")

def cmd_pci_isolate(devices):
    if pci_device_list_check(devices):
        return 1
    dev_tree = pci_tree_get()
    for dev in devices:
        pci_dev_isolate(pci_device_get(dev), True, "vfio-pci")

def cmd_pci_recover(devices):
    if not devices:
        devices = pci_device_list_get_by_driver(whitelist=["vfio-pci",
                                                           "pci-stub"])
        devices = map(lambda x: x["addr"], devices)
        if not devices:
            print("No isolated/stubbed device to recover")
            return 1
        print("Recovering all isolated/stubbed devices:")
        for dev in devices:
            print("  %s" % dev)
    if pci_device_list_check(devices):
        return 1
    dev_tree = pci_tree_get()
    for dev in devices:
        pci_dev_isolate(pci_device_get(dev), False)

def cmd_pci_routine(args):
    if len(args) <= 1:
        return pci_cmd_help()

    pci_cmd_list = {
        "isolate": cmd_pci_isolate,
        "recover": cmd_pci_recover,
        "status": cmd_pci_status,
        "stub": cmd_pci_stub,
    }

    subcmd = args[1]
    if subcmd not in pci_cmd_list:
        print("Unknown PCI subcommand: %s" % subcmd)
        return pci_cmd_help()

    return pci_cmd_list[subcmd](args[2:])

bz_url_search = "https://bugzilla.redhat.com/buglist.cgi?bug_status=NEW&bug_status=ASSIGNED&email1=peterx%40redhat.com&emailassigned_to1=1&emailtype1=substring&query_format=advanced"
bz_url_prefix = "https://bugzilla.redhat.com/show_bug.cgi?id="
org_bz_format = "** TODO [#B] %s\n   %s"

def cmd_list_bz(args):
    html_doc = requests.get(bz_url_search).text
    soup = BeautifulSoup(html_doc, 'html.parser')
    for link in soup.find_all("a"):
        if not bz_fmt.match(str(link.string)):
            continue
        bz_url = bz_url_prefix + link.string
        bz_doc = requests.get(bz_url).text
        bz = BeautifulSoup(bz_doc, 'html.parser')
        print(org_bz_format % (bz.title.string, bz_url))

def cmd_statistics(args):
    num_list = args[1:]
    if len(num_list) == 0:
        print("Please input a list of numbers for statistics.")
        return -1
    num_list = list(map(float, num_list))
    total = average = 0
    for i in num_list:
        total += i
    average = total / len(num_list)
    pos_max = neg_max = 0
    for i in num_list:
        diff = i - average
        if diff >= 0 and diff > pos_max:
            pos_max = diff
        elif diff < 0 and diff < neg_max:
            neg_max = diff
    diff = max(pos_max, -neg_max)
    diff = diff / average * 100
    print("Average: %.2f (+-%.2f%%)" % (average, diff))
    return 0

def cmd_numa_info(args):
    data = open("/proc/zoneinfo").read().split("\n")
    node_id = -1
    total = {}
    psize = resource.getpagesize()
    for line in data:
        if line.startswith("Node "):
            zone = line
        elif line.startswith("  pages free"):
            free_pages = int(line.replace("  pages free", "").strip())
            total[zone] = free_pages

    print("Free spaces for each zone:")
    for zone in total:
        space = int(total[zone] * psize / 1024 / 1024)
        if space:
            print("%s: %s (MB)" % (zone, space))
    return 0

CMD_HANDLERS = {
    "status-report": {
        "desc": "Generate status report",
        "handler": cmd_generate_status_report,
    },
    "pci": {
        "desc": "PCI related helper operations",
        "handler": cmd_pci_routine,
    },
    "list-bz": {
        "desc": "List RH BZs",
        "handler": cmd_list_bz,
    },
    "statistics": {
        "desc": "Do basic statistical calculation on an array",
        "handler": cmd_statistics,
    },
    "numa-info": {
        "desc": "Print numa/zone information",
        "handler": cmd_numa_info,
    },
}

def print_help():
    print("usage: %s <cmd>" % PROG_NAME)
    print()
    print("Command list:")
    print()
    for cmd in CMD_HANDLERS:
        print("%s - %s" % (cmd, CMD_HANDLERS[cmd]["desc"]))

if len(sys.argv) <= 1:
    print_help()
    exit(1)

cmd = sys.argv[1]

if cmd not in CMD_HANDLERS:
    print("Command %s unknown" % cmd)
    exit(1)

exit(CMD_HANDLERS[cmd]["handler"](sys.argv[1:]))