257 lines
6.3 KiB
Go
257 lines
6.3 KiB
Go
package tuimode
|
|
|
|
// A simple example demonstrating the use of multiple text input components
|
|
// from the Bubbles component library.
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"fmt"
|
|
"io"
|
|
"os/exec"
|
|
"regexp"
|
|
"strings"
|
|
|
|
"github.com/HikariKnight/ls-iommu/pkg/errorcheck"
|
|
"github.com/charmbracelet/bubbles/list"
|
|
tea "github.com/charmbracelet/bubbletea"
|
|
"github.com/charmbracelet/lipgloss"
|
|
)
|
|
|
|
var (
|
|
docStyle = lipgloss.NewStyle().Margin(2, 2)
|
|
titleStyle = lipgloss.NewStyle().
|
|
Background(lipgloss.Color("#5F5FD7")).
|
|
Foreground(lipgloss.Color("#FFFFFF"))
|
|
helpStyle = lipgloss.NewStyle().
|
|
Foreground(lipgloss.Color(241))
|
|
listStyle = lipgloss.NewStyle().
|
|
BorderStyle(lipgloss.RoundedBorder())
|
|
)
|
|
|
|
type status int
|
|
|
|
const (
|
|
GPUS status = iota
|
|
GPU_GROUP
|
|
USB
|
|
USB_GROUP
|
|
)
|
|
|
|
type item struct {
|
|
title, desc string
|
|
}
|
|
|
|
func (i item) Title() string { return i.title }
|
|
func (i item) Description() string { return i.desc }
|
|
func (i item) FilterValue() string { return i.title }
|
|
|
|
type model struct {
|
|
fetched []bool
|
|
lists []list.Model
|
|
loaded bool
|
|
focused status
|
|
width int
|
|
height int
|
|
}
|
|
|
|
func (m *model) initLists(width, height int) {
|
|
defaultList := list.New([]list.Item{}, list.NewDefaultDelegate(), width, height/2)
|
|
|
|
// Disable features we wont need
|
|
defaultList.SetShowTitle(false)
|
|
defaultList.SetFilteringEnabled(false)
|
|
defaultList.SetSize(width, height)
|
|
|
|
// Add height and width to our model so we can use it later
|
|
m.width = width
|
|
m.height = height
|
|
|
|
m.lists = []list.Model{defaultList, defaultList, defaultList, defaultList}
|
|
m.fetched = []bool{false, false, false, false}
|
|
m.focused = GPUS
|
|
|
|
// Init GPU list
|
|
//m.lists[GPUS].Title = "Select a GPU to check the IOMMU groups of"
|
|
items := GetIOMMU("-g", "-F", "name,device_id,optional_revision")
|
|
m.lists[GPUS].SetShowTitle(false)
|
|
m.lists[GPUS].SetItems(items)
|
|
m.fetched[GPUS] = true
|
|
|
|
m.lists[GPU_GROUP].Title = ""
|
|
m.lists[GPU_GROUP].SetItems(items)
|
|
|
|
// Init USB Controller list
|
|
items = GetIOMMU("-u", "-F", "name,device_id,optional_revision")
|
|
m.lists[USB].SetItems(items)
|
|
m.fetched[USB] = true
|
|
|
|
m.lists[USB_GROUP].Title = ""
|
|
m.lists[USB_GROUP].SetItems(items)
|
|
}
|
|
|
|
func (m model) Init() tea.Cmd {
|
|
return nil
|
|
}
|
|
|
|
// This function processes the enter event
|
|
func (m *model) processSelection() {
|
|
switch m.focused {
|
|
case GPUS:
|
|
// Gets the selected item
|
|
selectedItem := m.lists[m.focused].SelectedItem()
|
|
|
|
// Gets the IOMMU group of the selected item
|
|
iommu_group_regex := regexp.MustCompile(`(\d{1,3})`)
|
|
iommu_group := iommu_group_regex.FindString(selectedItem.(item).desc)
|
|
|
|
items := GetIOMMU("-gr", "-i", iommu_group, "-F", "name,device_id,optional_revision")
|
|
m.lists[GPU_GROUP].SetItems(items)
|
|
|
|
// Adjust height to correct for a bigger title
|
|
m.lists[GPU_GROUP].SetSize(m.width, m.height-1)
|
|
|
|
// Change focus to next index
|
|
m.focused++
|
|
|
|
case GPU_GROUP:
|
|
// Gets the selected item
|
|
/*selectedItem := m.lists[m.focused].SelectedItem()
|
|
|
|
// Gets the IOMMU group of the selected item
|
|
iommu_group_regex := regexp.MustCompile(`(\d{1,3})`)
|
|
iommu_group := iommu_group_regex.FindString(selectedItem.(item).desc)
|
|
|
|
items := GetIOMMU("-gr", "-i", iommu_group, "--id")*/
|
|
m.focused++
|
|
case USB:
|
|
// Gets the selected item
|
|
selectedItem := m.lists[m.focused].SelectedItem()
|
|
|
|
// Gets the IOMMU group of the selected item
|
|
iommu_group_regex := regexp.MustCompile(`(\d{1,3})`)
|
|
iommu_group := iommu_group_regex.FindString(selectedItem.(item).desc)
|
|
|
|
items := GetIOMMU("-ur", "-i", iommu_group, "-F", "name,device_id,optional_revision")
|
|
m.lists[USB_GROUP].SetItems(items)
|
|
|
|
// Adjust height to correct for a bigger title
|
|
m.lists[USB_GROUP].SetSize(m.width, m.height-1)
|
|
|
|
// Change focus to next index
|
|
m.focused++
|
|
}
|
|
}
|
|
|
|
func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
|
|
switch msg := msg.(type) {
|
|
case tea.KeyMsg:
|
|
switch msg.String() {
|
|
case "ctrl+c", "q":
|
|
return m, tea.Quit
|
|
|
|
case "enter":
|
|
if m.loaded {
|
|
m.processSelection()
|
|
}
|
|
case "ctrl+z", "backspace":
|
|
if m.focused > 0 {
|
|
m.focused--
|
|
return m, nil
|
|
} else {
|
|
return m, tea.Quit
|
|
}
|
|
}
|
|
case tea.WindowSizeMsg:
|
|
if !m.loaded {
|
|
// Get the terminal frame size
|
|
h, v := docStyle.GetFrameSize()
|
|
|
|
// Initialize the static lists and make sure the content
|
|
// does not extend past the screen
|
|
m.initLists(msg.Width-h, msg.Height-v)
|
|
|
|
// Set model loaded to true
|
|
m.loaded = true
|
|
}
|
|
}
|
|
|
|
var cmd tea.Cmd
|
|
m.lists[m.focused], cmd = m.lists[m.focused].Update(msg)
|
|
return m, cmd
|
|
}
|
|
|
|
func (m model) View() string {
|
|
if m.loaded {
|
|
title := ""
|
|
switch m.focused {
|
|
case GPUS:
|
|
title = " Select a GPU to check the IOMMU groups of"
|
|
|
|
case GPU_GROUP:
|
|
title = fmt.Sprint(
|
|
" Press ENTER/RETURN to set up all these devices for passthrough.\n",
|
|
" This list should only contain items related to your GPU.",
|
|
)
|
|
|
|
case USB:
|
|
title = " [OPTIONAL]: Select a USB Controller to check the IOMMU groups of"
|
|
|
|
case USB_GROUP:
|
|
title = fmt.Sprint(
|
|
" Press ENTER/RETURN to set up all these devices for passthrough.\n",
|
|
" This list should only contain the USB controller you want to use.",
|
|
)
|
|
}
|
|
return lipgloss.JoinVertical(lipgloss.Left, listStyle.SetString(fmt.Sprintf("%s\n", titleStyle.Render(title))).Render(m.lists[m.focused].View()))
|
|
} else {
|
|
return "Loading..."
|
|
}
|
|
}
|
|
|
|
func NewModel() *model {
|
|
// Create a blank model and return it
|
|
return &model{}
|
|
}
|
|
|
|
// This is where we build everything
|
|
func App() {
|
|
// Make a blank model to keep our state in
|
|
m := NewModel()
|
|
|
|
// Start the program with the model
|
|
p := tea.NewProgram(m, tea.WithAltScreen())
|
|
_, err := p.Run()
|
|
errorcheck.ErrorCheck(err, "Failed to initialize UI")
|
|
}
|
|
|
|
func GetIOMMU(args ...string) []list.Item {
|
|
var stdout, stderr bytes.Buffer
|
|
|
|
// Configure the ls-iommu command
|
|
cmd := exec.Command("utils/ls-iommu", args...)
|
|
cmd.Stderr = &stderr
|
|
cmd.Stdout = &stdout
|
|
|
|
// Execute the command
|
|
err := cmd.Run()
|
|
|
|
// If ls-iommu returns an error then IOMMU is disabled
|
|
errorcheck.ErrorCheck(err, "IOMMU disabled in either UEFI/BIOS or in bootloader!")
|
|
|
|
// Read the output
|
|
items := []list.Item{}
|
|
output, _ := io.ReadAll(&stdout)
|
|
|
|
// Parse the output line by line
|
|
scanner := bufio.NewScanner(strings.NewReader(string(output)))
|
|
for scanner.Scan() {
|
|
// Get the current line and split by :
|
|
objects := strings.Split(scanner.Text(), ": ")
|
|
// Write the objects into the list
|
|
items = append(items, item{title: objects[1], desc: objects[0]})
|
|
}
|
|
|
|
// Return our list of items
|
|
return items
|
|
}
|