usbguard "gui"

usbg

#!/usr/bin/env python3
import os
import re
import sys
import termios
import tty
import select
import shutil
import json
import subprocess

DEVICE_LIST_RE = r"(\d+):\s(allow|block)\sid\s([0-9a-f]{4}:[0-9a-f]{4})\sserial\s\"([^\"]*)\"\sname\s\"([^\"]*)\"\shash\s\"([^\"]+)\"\sparent-hash\s\"([^\"]+)\"\svia-port\s\"([^\"]+)\"\swith-interface\s(.*?)(?:\swith-connect-type\s\"[^\"]*\")?$"

# ANSI escape codes
UP = "\x1b[A"
DOWN = "\x1b[B"
SPACE = " "
ESCAPE = "\x1b"
CLEAR_SCREEN = "\x1b[2J\x1b[H"
CURSOR_UP = "\x1b[1A"
CURSOR_DOWN = "\x1b[1B"
CURSOR_SAVE = "\x1b[s"
CURSOR_RESTORE = "\x1b[u"
CURSOR_HIDE = "\x1b[?25l"
CURSOR_SHOW = "\x1b[?25h"

# Colors
RED = "\x1b[31m"
GREEN = "\x1b[32m"
RESET = "\x1b[0m"
WHITE_BG = "\x1b[47m"
BLACK_FG = "\x1b[30m"

# Check if terminal supports Unicode/emoji
if sys.stdout.isatty() and shutil.which("tput"):
    try:
        colors = int(subprocess.check_output(["tput", "colors"]).decode().strip())
        if colors >= 256:
            # Terminal supports colors and likely supports Unicode
            ALLOWED = "✅"
            BLOCKED = "❌"
        else:
            # Fallback to ASCII characters
            ALLOWED = " ✓"
            BLOCKED = " X"
    except (subprocess.SubprocessError, ValueError):
        # tput failed or couldn't parse output
        ALLOWED = " ✓"
        BLOCKED = " X"
else:
    # Not a terminal or tput not available, use ASCII
    ALLOWED = " ✓"
    BLOCKED = " X"

# Box drawing characters
BOX_TOP_LEFT = "┌"
BOX_TOP_RIGHT = "┐"
BOX_BOTTOM_LEFT = "└"
BOX_BOTTOM_RIGHT = "┘"
BOX_HORIZONTAL = "─"
BOX_VERTICAL = "│"
BOX_T_DOWN = "┬"
BOX_T_UP = "┴"
BOX_T_RIGHT = "├"
BOX_T_LEFT = "┤"
BOX_CROSS = "┼"

# Helper to strip ANSI codes for width calculation
ANSI_ESCAPE = re.compile(r"\x1b\[[0-9;]*m")


def strip_ansi(s):
    return ANSI_ESCAPE.sub("", s)


def get_char():
    """Get a single character from stdin without requiring Enter key."""
    fd = sys.stdin.fileno()
    old_settings = termios.tcgetattr(fd)
    try:
        tty.setraw(sys.stdin.fileno())
        ch = sys.stdin.read(1)
        if ch == "\x1b":
            ch2 = sys.stdin.read(1)
            ch3 = sys.stdin.read(1)
            if ch2 == "[":
                if ch3 == "A":
                    return UP
                elif ch3 == "B":
                    return DOWN
        return ch
    finally:
        termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)


def get_usb_info(vendor_id, product_id):
    """Look up vendor and device information from USB IDs."""
    try:
        # Fallback to usb.ids database
        usb_ids_path = "/usr/share/misc/usb.ids"
        if not os.path.exists(usb_ids_path):
            print(
                f"Warning: {usb_ids_path} not found. USB vendor/device info will be limited."
            )
            return "Unknown Vendor", "Unknown Device"
        try:
            with open(usb_ids_path, "r", encoding="utf-8", errors="ignore") as f:
                lines = f.readlines()
            vendor_name = None
            device_name = None
            in_vendor = False

            for line in lines:
                if line.startswith("#") or not line.strip():
                    continue
                # Use a simpler regex pattern for vendor lines
                if re.match(r"^[0-9a-fA-F]{4}\s", line):
                    current_vendor_id = line[:4].lower()
                    if current_vendor_id == vendor_id.lower():
                        vendor_name = line[4:].strip()
                        in_vendor = True
                    else:
                        in_vendor = False
                elif in_vendor and line.startswith("\t"):
                    # Use a simpler regex pattern for product lines
                    if re.match(r"\t[0-9a-fA-F]{4}\s", line):
                        current_product_id = line[1:5].lower()
                        if current_product_id == product_id.lower():
                            device_name = line[6:].strip()
                            break
            if not vendor_name:
                vendor_name = "Unknown Vendor"
            if not device_name:
                device_name = "Unknown Device"
            return vendor_name, device_name
        except Exception as e:
            print(f"Warning: Failed to parse {usb_ids_path}: {e}")
            return "Unknown Vendor", "Unknown Device"
    except Exception as e:
        print(f"Warning: Exception in get_usb_info: {e}")
        return "Unknown Vendor", "Unknown Device"


def get_device_list():
    stream = os.popen("usbguard list-devices")
    raw_output = stream.read().strip()

    # Join wrapped lines
    lines = []
    current_line = ""
    for line in raw_output.splitlines():
        line = line.strip()
        if not line:  # Skip empty lines
            continue
        if (
            line[0].isdigit() and ": " in line
        ):  # New device entry starts with number and colon
            if current_line:
                lines.append(current_line)
            current_line = line
        else:  # Continuation of previous line
            current_line += " " + line

    if current_line:
        lines.append(current_line)

    devices = []
    for line in lines:
        # Extract device ID and status first
        id_match = re.match(r"(\d+):\s(allow|block)", line)
        if not id_match:
            continue

        # Extract the rest of the information
        rest = line[id_match.end() :]
        device = {}
        device["usb_id"] = int(id_match.group(1))
        device["usb_status"] = id_match.group(2)

        # Extract other fields using simpler patterns
        id_str_match = re.search(r"id\s([0-9a-f]{4}:[0-9a-f]{4})", rest)
        serial_match = re.search(r'serial\s"([^"]*)"', rest)
        name_match = re.search(r'name\s"([^"]*)"', rest)
        hash_match = re.search(r'hash\s"([^"]+)"', rest)
        parent_hash_match = re.search(r'parent-hash\s"([^\"]+)"', rest)
        port_match = re.search(r'via-port\s"([^\"]+)"', rest)
        interface_match = re.search(
            r"with-interface\s(.*?)(?:\swith-connect-type|$)", rest
        )

        if id_str_match:
            device["usb_id_str"] = id_str_match.group(1)
            # Get vendor and device info
            vendor_id, product_id = device["usb_id_str"].split(":")
            vendor_name, device_name = get_usb_info(vendor_id, product_id)
            device["vendor_name"] = vendor_name or "Unknown Vendor"
            device["device_name"] = device_name or "Unknown Device"
        else:
            continue

        device["serial"] = serial_match.group(1) if serial_match else ""
        # Use device name from USB IDs if usbguard output has no name
        if name_match and name_match.group(1).strip():
            device["usb_name"] = name_match.group(1).strip()
        else:
            device["usb_name"] = device[
                "device_name"
            ]  # Use the device name from USB IDs lookup
        device["hash"] = hash_match.group(1) if hash_match else ""
        device["parent_hash"] = parent_hash_match.group(1) if parent_hash_match else ""
        device["port"] = port_match.group(1) if port_match else ""
        device["interfaces"] = (
            interface_match.group(1).strip() if interface_match else ""
        )

        devices.append(device)

    return devices, [device["usb_id"] for device in devices]


def find_device_index(devices, device_id):
    """Find the index of a device by its ID."""
    for i, device in enumerate(devices):
        if device["usb_id"] == device_id:
            return i
    return 0  # Default to first device if not found


def main():
    print(CURSOR_HIDE)  # Hide cursor
    try:
        selected_device_id = None  # Track selected device by ID instead of index
        show_help = False
        tree_view = True  # Default to tree view

        while True:
            devices, _ = get_device_list()
            if not devices:
                print("No USB devices found")
                break

            # Track the order of devices as they are printed
            selected_index = 0
            if selected_device_id is not None:
                # We'll set selected_index after getting printed_order
                pass

            def print_device_list(
                devices, selected_index, show_help=False, tree_view=True
            ):
                print(CLEAR_SCREEN)

                # Tree drawing characters
                TREE_VERTICAL = "│"
                TREE_BRANCH = "├"
                TREE_LAST = "└"
                TREE_HORIZONTAL = "─"
                TREE_SPACE = " "

                # Group devices by their parent hash
                device_tree = {}
                root_devices = []
                for device in devices:
                    parent_hash = device.get("parent_hash", "")
                    if (
                        parent_hash == "0000000000000000000000000000000000000000000000"
                        or not parent_hash
                    ):
                        root_devices.append(device)
                    else:
                        if parent_hash not in device_tree:
                            device_tree[parent_hash] = []
                        device_tree[parent_hash].append(device)

                printed_devices = set()
                printed_order = []
                row_data = (
                    []
                )  # Will hold tuples: (level, is_last, prefix, device, device_name)
                max_depth = 0
                max_name_len = len("Device Name")
                max_vendor_len = len("Vendor")
                max_device_len = len("Device")

                def collect_rows(device, level=0, is_last=True, prefix=""):
                    nonlocal max_depth, max_name_len, max_vendor_len, max_device_len
                    if device["usb_id"] in printed_devices:
                        return
                    printed_devices.add(device["usb_id"])
                    printed_order.append(device["usb_id"])
                    # Calculate indentation
                    indent = prefix
                    if level > 0:
                        indent += TREE_LAST if is_last else TREE_BRANCH
                        indent += TREE_HORIZONTAL * 2 + " "
                    device_name = f"{indent}{device['usb_name']}"
                    row_data.append((level, is_last, prefix, device, device_name))
                    max_depth = max(max_depth, level)
                    max_name_len = max(max_name_len, len(device_name))
                    max_vendor_len = max(max_vendor_len, len(device["vendor_name"]))
                    max_device_len = max(max_device_len, len(device["device_name"]))
                    # Children
                    children = device_tree.get(device.get("hash", ""), [])
                    for i, child in enumerate(children):
                        new_prefix = prefix + (
                            TREE_SPACE * 3
                            if is_last
                            else TREE_VERTICAL + TREE_SPACE * 2
                        )
                        collect_rows(
                            child, level + 1, i == len(children) - 1, new_prefix
                        )

                if tree_view:
                    if not root_devices:
                        root_devices = devices
                    for i, device in enumerate(root_devices):
                        collect_rows(device, 0, i == len(root_devices) - 1)
                else:
                    for idx, device in enumerate(devices):
                        row_data.append((0, True, "", device, device["usb_name"]))
                        printed_order.append(device["usb_id"])
                        max_name_len = max(max_name_len, len(device["usb_name"]))
                        max_vendor_len = max(max_vendor_len, len(device["vendor_name"]))
                        max_device_len = max(max_device_len, len(device["device_name"]))

                # Now calculate column widths
                status_width = 8
                name_width = max_name_len
                vendor_width = max_vendor_len
                device_width = max_device_len
                id_width = max(
                    len("Device ID"),
                    max(len(device["usb_id_str"]) for device in devices),
                )

                status_total = status_width + 3
                name_total = name_width + 2
                vendor_total = vendor_width + 2
                device_total = device_width + 2
                id_total = id_width + 2

                # Print table header
                header = (
                    f"{BOX_TOP_LEFT}{BOX_HORIZONTAL * status_total}"
                    f"{BOX_T_DOWN}{BOX_HORIZONTAL * name_total}"
                    f"{BOX_T_DOWN}{BOX_HORIZONTAL * vendor_total}"
                    f"{BOX_T_DOWN}{BOX_HORIZONTAL * device_total}"
                    f"{BOX_T_DOWN}{BOX_HORIZONTAL * id_total}"
                    f"{BOX_TOP_RIGHT}"
                )
                print(header)
                print(
                    f"{BOX_VERTICAL} {'Status':<{(status_width + 1)}} {BOX_VERTICAL} "
                    f"{'Device Name':<{name_width}} {BOX_VERTICAL} "
                    f"{'Vendor':<{vendor_width}} {BOX_VERTICAL} "
                    f"{'Device':<{device_width}} {BOX_VERTICAL} "
                    f"{'Device ID':<{id_width}} {BOX_VERTICAL}"
                )
                separator = (
                    f"{BOX_T_RIGHT}{BOX_HORIZONTAL * status_total}"
                    f"{BOX_CROSS}{BOX_HORIZONTAL * name_total}"
                    f"{BOX_CROSS}{BOX_HORIZONTAL * vendor_total}"
                    f"{BOX_CROSS}{BOX_HORIZONTAL * device_total}"
                    f"{BOX_CROSS}{BOX_HORIZONTAL * id_total}"
                    f"{BOX_T_LEFT}"
                )
                print(separator)

                # Print rows
                for idx, (level, is_last, prefix, device, device_name) in enumerate(
                    row_data
                ):
                    if device["usb_status"] == "allow":
                        status_colored = ALLOWED
                    else:
                        status_colored = BLOCKED
                    total_pad = status_width - 1
                    left_pad = total_pad // 2
                    right_pad = total_pad - left_pad
                    status_cell = f"{' ' * left_pad}{status_colored}{' ' * right_pad}"
                    if printed_order.index(device["usb_id"]) == selected_index:
                        print(
                            f"{BOX_VERTICAL}{WHITE_BG}{BLACK_FG} {status_cell} {BOX_VERTICAL} "
                            f"{device_name:<{name_width}} {BOX_VERTICAL} "
                            f"{device['vendor_name']:<{vendor_width}} {BOX_VERTICAL} "
                            f"{device['device_name']:<{device_width}} {BOX_VERTICAL} "
                            f"{device['usb_id_str']:<{id_width}} {RESET}{BOX_VERTICAL}"
                        )
                    else:
                        print(
                            f"{BOX_VERTICAL} {status_cell} {BOX_VERTICAL} "
                            f"{device_name:<{name_width}} {BOX_VERTICAL} "
                            f"{device['vendor_name']:<{vendor_width}} {BOX_VERTICAL} "
                            f"{device['device_name']:<{device_width}} {BOX_VERTICAL} "
                            f"{device['usb_id_str']:<{id_width}} {BOX_VERTICAL}"
                        )

                # Print table footer
                footer = (
                    f"{BOX_BOTTOM_LEFT}{BOX_HORIZONTAL * status_total}"
                    f"{BOX_T_UP}{BOX_HORIZONTAL * name_total}"
                    f"{BOX_T_UP}{BOX_HORIZONTAL * vendor_total}"
                    f"{BOX_T_UP}{BOX_HORIZONTAL * device_total}"
                    f"{BOX_T_UP}{BOX_HORIZONTAL * id_total}"
                    f"{BOX_BOTTOM_RIGHT}"
                )
                print(footer)

                # Print help indicator and help text below the table
                print(f"Press 'h' for help")
                if show_help:
                    print("\nControls:")
                    print("↑/↓ arrows - Navigate")
                    print("SPACE     - Toggle device")
                    print("r         - Refresh list")
                    print("h         - Toggle help")
                    print("t         - Toggle tree/flat view")
                    print("q         - Quit")

                return printed_order

            printed_order = print_device_list(
                devices, selected_index, show_help, tree_view
            )

            # Set selected_index to match selected_device_id in the current printed_order
            if selected_device_id is None:
                selected_index = 0
                selected_device_id = printed_order[0] if printed_order else None
            else:
                try:
                    selected_index = printed_order.index(selected_device_id)
                except ValueError:
                    selected_index = 0
                    selected_device_id = printed_order[0] if printed_order else None

            # Reprint with correct selected_index
            printed_order = print_device_list(
                devices, selected_index, show_help, tree_view
            )
            selected_device_id = (
                printed_order[selected_index] if printed_order else None
            )

            while True:
                char = get_char()

                if char == "q":
                    print(CURSOR_SHOW)  # Show cursor before exit
                    return

                elif char == "r":
                    break  # Break inner loop to refresh device list

                elif char == "h":
                    show_help = not show_help
                    printed_order = print_device_list(
                        devices, selected_index, show_help, tree_view
                    )

                elif char == "t":
                    tree_view = not tree_view
                    # After toggling, update selected_index to match selected_device_id in new order
                    try:
                        selected_index = printed_order.index(selected_device_id)
                    except ValueError:
                        selected_index = 0
                        selected_device_id = printed_order[0] if printed_order else None
                    printed_order = print_device_list(
                        devices, selected_index, show_help, tree_view
                    )

                elif char == UP:
                    if selected_index > 0:
                        selected_index -= 1
                        selected_device_id = printed_order[selected_index]
                        printed_order = print_device_list(
                            devices, selected_index, show_help, tree_view
                        )

                elif char == DOWN:
                    if selected_index < len(printed_order) - 1:
                        selected_index += 1
                        selected_device_id = printed_order[selected_index]
                        printed_order = print_device_list(
                            devices, selected_index, show_help, tree_view
                        )

                elif char == SPACE:
                    device_id = printed_order[selected_index]
                    if toggle_device(device_id, devices):
                        printed_order = print_device_list(
                            devices, selected_index, show_help, tree_view
                        )
                    break  # Refresh the device list after toggle

    finally:
        print(CURSOR_SHOW)  # Ensure cursor is shown before exit


def toggle_device(device_id, devices):
    for item in devices:
        if item["usb_id"] == device_id:
            if item["usb_status"] == "allow":
                os.popen("usbguard block-device {}".format(device_id))
                item["usb_status"] = "block"
            else:
                os.popen("usbguard allow-device {}".format(device_id))
                item["usb_status"] = "allow"
            return True
    return False


if __name__ == "__main__":
    main()