Implement TUI for managing suspicious artifacts (FilePicker, TypePicker, ListView, and DetailView)
Introduce Bubble Tea-based terminal UI to manage suspicious artifact findings, including file selection, type filtering, list view, and details.
This commit is contained in:
Evan Hosinski
@@ -0,0 +1,241 @@
|
||||
package tui
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
|
||||
"rmm-hunter/internal/suspicious"
|
||||
|
||||
tea "github.com/charmbracelet/bubbletea"
|
||||
"github.com/charmbracelet/lipgloss"
|
||||
)
|
||||
|
||||
type screen int
|
||||
|
||||
const (
|
||||
screenFilePicker screen = iota
|
||||
screenTypePicker
|
||||
screenList
|
||||
screenDetail
|
||||
screenError
|
||||
)
|
||||
|
||||
type AppModel struct {
|
||||
current screen
|
||||
filePick FilePickerModel
|
||||
typePick TypePickerModel
|
||||
listView ListViewModel
|
||||
detail DetailViewModel
|
||||
err error
|
||||
selected string
|
||||
data suspicious.Suspicious
|
||||
width int
|
||||
height int
|
||||
}
|
||||
|
||||
func NewApp() AppModel {
|
||||
return AppModel{
|
||||
current: screenFilePicker,
|
||||
filePick: NewFilePicker(),
|
||||
typePick: NewTypePicker(),
|
||||
}
|
||||
}
|
||||
|
||||
func (m AppModel) Init() tea.Cmd { return m.filePick.Init() }
|
||||
|
||||
func (m AppModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
|
||||
// remember the latest terminal size so we can size new screens
|
||||
if ws, ok := msg.(tea.WindowSizeMsg); ok {
|
||||
m.width, m.height = ws.Width, ws.Height
|
||||
}
|
||||
|
||||
switch m.current {
|
||||
case screenFilePicker:
|
||||
var cmd tea.Cmd
|
||||
var tm tea.Model
|
||||
tm, cmd = m.filePick.Update(msg)
|
||||
if fp, ok := tm.(FilePickerModel); ok {
|
||||
m.filePick = fp
|
||||
}
|
||||
switch v := msg.(type) {
|
||||
case FileSelectedMsg:
|
||||
if err := m.loadSelectedFile(v.Path); err != nil {
|
||||
m.err = err
|
||||
m.current = screenError
|
||||
return m, nil
|
||||
}
|
||||
m.current = screenTypePicker
|
||||
return m, nil
|
||||
}
|
||||
return m, cmd
|
||||
|
||||
case screenTypePicker:
|
||||
var cmd tea.Cmd
|
||||
var tm tea.Model
|
||||
tm, cmd = m.typePick.Update(msg)
|
||||
if tp, ok := tm.(TypePickerModel); ok {
|
||||
m.typePick = tp
|
||||
}
|
||||
switch v := msg.(type) {
|
||||
case BackMsg:
|
||||
m.current = screenFilePicker
|
||||
return m, nil
|
||||
case SelectedTypeMsg:
|
||||
m.selected = v.Type
|
||||
m.listView = NewListView(v.Type, m.data, m.width, m.height)
|
||||
m.current = screenList
|
||||
return m, nil
|
||||
}
|
||||
return m, cmd
|
||||
|
||||
case screenList:
|
||||
var cmd tea.Cmd
|
||||
var tm tea.Model
|
||||
tm, cmd = m.listView.Update(msg)
|
||||
if lv, ok := tm.(ListViewModel); ok {
|
||||
m.listView = lv
|
||||
}
|
||||
switch v := msg.(type) {
|
||||
case BackMsg:
|
||||
m.current = screenTypePicker
|
||||
return m, nil
|
||||
case ListSelectedMsg:
|
||||
m.detail = NewDetailView(v.TypeKey, v.Index, m.data)
|
||||
m.current = screenDetail
|
||||
return m, nil
|
||||
}
|
||||
return m, cmd
|
||||
|
||||
case screenDetail:
|
||||
var cmd tea.Cmd
|
||||
var tm tea.Model
|
||||
tm, cmd = m.detail.Update(msg)
|
||||
if dv, ok := tm.(DetailViewModel); ok {
|
||||
m.detail = dv
|
||||
}
|
||||
switch v := msg.(type) {
|
||||
case BackMsg:
|
||||
m.current = screenList
|
||||
return m, nil
|
||||
case RequestEliminateMsg:
|
||||
if err := m.performEliminate(v.TypeKey, v.Index); err != nil {
|
||||
m.detail.modalErr = err.Error()
|
||||
return m, nil
|
||||
}
|
||||
// success -> rebuild list and go back
|
||||
m.listView = NewListView(m.selected, m.data, m.width, m.height)
|
||||
m.current = screenList
|
||||
return m, nil
|
||||
}
|
||||
return m, cmd
|
||||
|
||||
case screenError:
|
||||
// Any key quits after error is shown
|
||||
if _, ok := msg.(tea.KeyMsg); ok {
|
||||
return m, tea.Quit
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
|
||||
func (m AppModel) View() string {
|
||||
switch m.current {
|
||||
case screenFilePicker:
|
||||
return m.filePick.View()
|
||||
case screenTypePicker:
|
||||
return m.typePick.View()
|
||||
case screenList:
|
||||
return m.listView.View()
|
||||
case screenDetail:
|
||||
return m.detail.View()
|
||||
case screenError:
|
||||
return lipgloss.NewStyle().Foreground(lipgloss.Color("203")).Render(fmt.Sprintf("Failed to load JSON: %v\nPress any key to exit.", m.err))
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
// performEliminate routes to placeholder eliminate functions and mutates data on success
|
||||
func (m *AppModel) performEliminate(typeKey string, idx int) error {
|
||||
switch typeKey {
|
||||
case "autoruns":
|
||||
ar := m.data.AutoRuns[idx]
|
||||
if err := EliminateAutoRun(ar); err != nil {
|
||||
return err
|
||||
}
|
||||
m.data.AutoRuns = append(m.data.AutoRuns[:idx], m.data.AutoRuns[idx+1:]...)
|
||||
case "binaries":
|
||||
b := m.data.Binaries[idx]
|
||||
if err := EliminateBinary(b); err != nil {
|
||||
return err
|
||||
}
|
||||
m.data.Binaries = append(m.data.Binaries[:idx], m.data.Binaries[idx+1:]...)
|
||||
case "connections":
|
||||
c := m.data.OutboundConnections[idx]
|
||||
if err := EliminateConnection(c); err != nil {
|
||||
return err
|
||||
}
|
||||
m.data.OutboundConnections = append(m.data.OutboundConnections[:idx], m.data.OutboundConnections[idx+1:]...)
|
||||
case "directories":
|
||||
d := m.data.Directories[idx]
|
||||
if err := EliminateDirectory(d); err != nil {
|
||||
return err
|
||||
}
|
||||
m.data.Directories = append(m.data.Directories[:idx], m.data.Directories[idx+1:]...)
|
||||
case "processes":
|
||||
p := m.data.Processes[idx]
|
||||
if err := EliminateProcess(p); err != nil {
|
||||
return err
|
||||
}
|
||||
m.data.Processes = append(m.data.Processes[:idx], m.data.Processes[idx+1:]...)
|
||||
case "scheduledTasks":
|
||||
t := m.data.ScheduledTasks[idx]
|
||||
if err := EliminateScheduledTask(*t); err != nil {
|
||||
return err
|
||||
}
|
||||
m.data.ScheduledTasks = append(m.data.ScheduledTasks[:idx], m.data.ScheduledTasks[idx+1:]...)
|
||||
case "services":
|
||||
s := m.data.Services[idx]
|
||||
if err := EliminateService(*s); err != nil {
|
||||
return err
|
||||
}
|
||||
m.data.Services = append(m.data.Services[:idx], m.data.Services[idx+1:]...)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// loadSelectedFile reads the JSON file and populates m.data
|
||||
func (m *AppModel) loadSelectedFile(path string) error {
|
||||
b, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// Support both wrapped report (with findings) and bare Suspicious JSON
|
||||
var envelope struct {
|
||||
Findings json.RawMessage `json:"findings"`
|
||||
}
|
||||
if err := json.Unmarshal(b, &envelope); err == nil && len(envelope.Findings) > 0 {
|
||||
var sus suspicious.Suspicious
|
||||
if err := json.Unmarshal(envelope.Findings, &sus); err != nil {
|
||||
return err
|
||||
}
|
||||
m.data = sus
|
||||
return nil
|
||||
}
|
||||
// Try bare suspicious structure
|
||||
var sus suspicious.Suspicious
|
||||
if err := json.Unmarshal(b, &sus); err != nil {
|
||||
return fmt.Errorf("no findings in report")
|
||||
}
|
||||
m.data = sus
|
||||
return nil
|
||||
}
|
||||
|
||||
// RunEliminateUI starts the Bubble Tea program for elimination UI
|
||||
func RunEliminateUI() error {
|
||||
p := tea.NewProgram(NewApp())
|
||||
_, err := p.Run()
|
||||
return err
|
||||
}
|
||||
Reference in New Issue
Block a user