quickpassthrough/internal/tuimode/tuimode.go
2023-12-26 19:20:28 +01:00

261 lines
6.4 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 = fmt.Sprintf(
" 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 = fmt.Sprint(
" [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
}