package tuimode // A simple example demonstrating the use of multiple text input components // from the Bubbles component library. import ( "bufio" "bytes" "fmt" "io" "os" "os/exec" "path/filepath" "regexp" "strings" "github.com/HikariKnight/ls-iommu/pkg/errorcheck" "github.com/HikariKnight/quickpassthrough/internal/configs" "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")). PaddingLeft(2).PaddingRight(2) helpStyle = lipgloss.NewStyle(). Foreground(lipgloss.Color(241)) listStyle = lipgloss.NewStyle(). PaddingLeft(2) choiceStyle = lipgloss.NewStyle().PaddingLeft(4) selectedChoiceStyle = lipgloss.NewStyle(). PaddingLeft(2). Foreground(lipgloss.Color("170")) dialogStyle = lipgloss.NewStyle(). PaddingLeft(2). Width(78) ) // Make a status type type status int // List item struct type item struct { title, desc string } // Functions needed for item struct func (i item) Title() string { return i.title } func (i item) Description() string { return i.desc } func (i item) FilterValue() string { return i.title } // Choice delegate (for our dialog boxes) type choiceDelegate struct{} func (d choiceDelegate) Height() int { return 1 } func (d choiceDelegate) Spacing() int { return 0 } func (d choiceDelegate) Update(msg tea.Msg, m *list.Model) tea.Cmd { return nil } func (d choiceDelegate) Render(w io.Writer, m list.Model, index int, listItem list.Item) { i, ok := listItem.(item) if !ok { return } str := fmt.Sprintf("%s", i.title) fn := choiceStyle.Render if index == m.Index() { fn = func(s ...string) string { return selectedChoiceStyle.Render("| " + strings.Join(s, " ")) } } fmt.Fprint(w, fn(str)) } // Main Model type model struct { fetched []bool lists []list.Model gpu_group string vbios_path string loaded bool focused status width int height int } // Consts used to navigate the main model const ( GPUS status = iota GPU_GROUP VBIOS USB USB_GROUP ) func (m *model) initLists(width, height int) { defaultList := list.New([]list.Item{}, list.NewDefaultDelegate(), 0, 10) choiceList := list.New([]list.Item{}, choiceDelegate{}, 0, 7) // Disable features we wont need defaultList.SetShowTitle(false) defaultList.SetFilteringEnabled(false) defaultList.SetSize(width, height) choiceList.SetShowTitle(false) choiceList.SetFilteringEnabled(false) // 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, choiceList, defaultList, defaultList, } m.fetched = []bool{ false, false, false, false, false, } m.focused = GPUS // Init GPU list //m.lists[GPUS].Title = "Select a GPU to check the IOMMU groups of" items := StringList2ListItem(GetIOMMU("-g", "-F", "name,device_id,optional_revision")) 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 = StringList2ListItem(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) items = []list.Item{ item{title: "OK"}, } m.lists[VBIOS].SetItems(items) //m.lists[TEST].SetSize(width, height) } func (m model) Init() tea.Cmd { return nil } // This function processes the enter event func (m *model) processSelection() { switch m.focused { case GPUS: configs.InitConfigs() // 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) // Add the gpu group to our model m.gpu_group = iommu_group items := StringList2ListItem(GetIOMMU("-gr", "-i", m.gpu_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: // Generate the VBIOS dumper script once the user has selected a GPU GenerateVBIOSDumper(*m) 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 := StringList2ListItem(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++ case USB_GROUP: m.focused++ case VBIOS: 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-2, msg.Height-2) // 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 = titleStyle.Render( "Select a GPU to check the IOMMU groups of", ) case GPU_GROUP: title = titleStyle.Render( 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 = titleStyle.Render( "[OPTIONAL]: Select a USB Controller to check the IOMMU groups of", ) case USB_GROUP: title = titleStyle.Render( 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.", ), ) case VBIOS: // Get the program directory exe, _ := os.Executable() scriptdir := filepath.Dir(exe) // If we are using go run use the working directory instead if strings.Contains(scriptdir, "/tmp/go-build") { scriptdir, _ = os.Getwd() } text := dialogStyle.Render( fmt.Sprint( "Based on your GPU selection, a vbios extraction script has been generated for your convenience.\n", "Passing a VBIOS rom to the card used for passthrough is required for some cards, but not all.\n", "Some cards also requires you to patch your VBIOS romfile, check online if this is neccessary for your card!\n", "The VBIOS will be read from:\n", "%s\n\n", "The script to extract the vbios has to be run as sudo and without a displaymanager running for proper dumping!\n", "\n", "You can run the script with:\n", "%s/utils/dump_vbios.sh", ), ) title = fmt.Sprintf(text, m.vbios_path, scriptdir) } //return listStyle.SetString(fmt.Sprintf("%s\n\n", title)).Render(m.lists[m.focused].View()) return lipgloss.JoinVertical(lipgloss.Left, fmt.Sprintf("%s\n%s\n", title, listStyle.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) []string { 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 var items []string output, _ := io.ReadAll(&stdout) // Parse the output line by line scanner := bufio.NewScanner(strings.NewReader(string(output))) for scanner.Scan() { // Write the objects into the list items = append(items, scanner.Text()) } // Return our list of items return items } func StringList2ListItem(stringList []string) []list.Item { // Make the []list.Item struct items := []list.Item{} // Parse the output line by line for _, v := range stringList { // Get the current line and split by : objects := strings.Split(v, ": ") // Write the objects into the list items = append(items, item{title: objects[1], desc: objects[0]}) } // Return our list of items return items } func GenerateVBIOSDumper(m model) { // Get the vbios path m.vbios_path = GetIOMMU("-g", "-i", m.gpu_group, "--rom")[0] // Get the config directories config := configs.GetConfigPaths() // Get the program directory exe, _ := os.Executable() scriptdir := filepath.Dir(exe) // If we are using go run use the working directory instead if strings.Contains(scriptdir, "/tmp/go-build") { scriptdir, _ = os.Getwd() } vbios_script_template := fmt.Sprint( "#!/bin/bash\n", "# THIS FILE IS AUTO GENERATED!\n", "# IF YOU HAVE CHANGED GPU, PLEASE RE-RUN QUICKPASSTHROUGH!\n", "echo 1 | sudo tee %s\n", "sudo bash -c \"cat %s\" > %s/%s/vfio_card.rom\n", "echo 0 | sudo tee %s\n", ) vbios_script := fmt.Sprintf( vbios_script_template, m.vbios_path, m.vbios_path, scriptdir, config.QUICKEMU, m.vbios_path, ) scriptfile, err := os.Create("utils/dump_vbios.sh") errorcheck.ErrorCheck(err, "Cannot create file \"utils/dump_vbios.sh\"") defer scriptfile.Close() // Make the script executable scriptfile.Chmod(0775) errorcheck.ErrorCheck(err, "Could not change permissions of \"utils/dump_vbios.sh\"") // Write the script scriptfile.WriteString(vbios_script) }