Refactor suspicious artifact data structures, enhance eliminated state tracking, and update UI rendering for eliminated items. Add JSON marshal/unmarshal support for Binary and Directory types.

2025-10-11 21:01:07 -04:00
parent bde1b23753
commit c9e2e8dff8
9 changed files with 409 additions and 88 deletions
@@ -23,23 +23,26 @@ const (
 )

 type AppModel struct {
-	current  screen
-	filePick FilePickerModel
-	typePick TypePickerModel
-	listView ListViewModel
-	detail   DetailViewModel
-	err      error
-	selected string
-	data     suspicious.Suspicious
-	width    int
-	height   int
+	current    screen
+	filePick   FilePickerModel
+	typePick   TypePickerModel
+	listView   ListViewModel
+	detail     DetailViewModel
+	err        error
+	selected   string
+	data       suspicious.Suspicious
+	width      int
+	height     int
+	eliminated map[string]map[int]bool // tracks eliminated items: typeKey -> index -> eliminated
+	filePath   string                  // path to the loaded JSON file
 }

 func NewApp() AppModel {
 	return AppModel{
-		current:  screenFilePicker,
-		filePick: NewFilePicker(),
-		typePick: NewTypePicker(),
+		current:    screenFilePicker,
+		filePick:   NewFilePicker(),
+		typePick:   NewTypePicker(),
+		eliminated: make(map[string]map[int]bool),
 	}
 }

@@ -66,6 +69,7 @@ func (m AppModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 				m.current = screenError
 				return m, nil
 			}
+			m.filePath = v.Path
 			m.current = screenTypePicker
 			return m, nil
 		}
@@ -84,7 +88,7 @@ func (m AppModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 			return m, nil
 		case SelectedTypeMsg:
 			m.selected = v.Type
-			m.listView = NewListView(v.Type, m.data, m.width, m.height)
+			m.listView = NewListView(v.Type, m.data, m.width, m.height, m.eliminated)
 			m.current = screenList
 			return m, nil
 		}
@@ -102,7 +106,7 @@ func (m AppModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 			m.current = screenTypePicker
 			return m, nil
 		case ListSelectedMsg:
-			m.detail = NewDetailView(v.TypeKey, v.Index, m.data)
+			m.detail = NewDetailView(v.TypeKey, v.Index, m.data, m.eliminated)
 			m.current = screenDetail
 			return m, nil
 		}
@@ -120,6 +124,12 @@ func (m AppModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 			m.current = screenList
 			return m, nil
 		case RequestEliminateMsg:
+			// Check if already eliminated
+			if m.eliminated[v.TypeKey] != nil && m.eliminated[v.TypeKey][v.Index] {
+				m.detail.modalWarn = "This item has already been eliminated"
+				return m, nil
+			}
+
 			if err := m.performEliminate(v.TypeKey, v.Index); err != nil {
 				var wb WarnBlock
 				if errors.As(err, &wb) {
@@ -129,8 +139,20 @@ func (m AppModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 				}
 				return m, nil
 			}
-			// success -> rebuild list and go back
-			m.listView = NewListView(m.selected, m.data, m.width, m.height)
+			// success -> mark as eliminated, save file, and rebuild list
+			if m.eliminated[v.TypeKey] == nil {
+				m.eliminated[v.TypeKey] = make(map[int]bool)
+			}
+			m.eliminated[v.TypeKey][v.Index] = true
+
+			// Save the updated data to file
+			if err := m.saveDataToFile(); err != nil {
+				m.detail.modalErr = fmt.Sprintf("Eliminated successfully but failed to save: %v", err)
+				return m, nil
+			}
+
+			m.detail = NewDetailView(v.TypeKey, v.Index, m.data, m.eliminated)
+			m.listView = NewListView(m.selected, m.data, m.width, m.height, m.eliminated)
 			m.current = screenList
 			return m, nil
 		}
@@ -163,7 +185,7 @@ func (m AppModel) View() string {
 	}
 }

-// performEliminate routes to placeholder eliminate functions and mutates data on success
+// performEliminate routes to placeholder eliminate functions without removing items from data
 func (m *AppModel) performEliminate(typeKey string, idx int) error {
 	switch typeKey {
 	case "autoruns":
@@ -171,49 +193,49 @@ func (m *AppModel) performEliminate(typeKey string, idx int) error {
 		if err := EliminateAutoRun(ar); err != nil {
 			return err
 		}
-		m.data.AutoRuns = append(m.data.AutoRuns[:idx], m.data.AutoRuns[idx+1:]...)
+		m.data.AutoRuns[idx].Eliminated = true
 	case "binaries":
 		b := m.data.Binaries[idx]
-		if err := CheckBinaryBlocked(b, m.data); err != nil {
+		if err := CheckBinaryBlocked(b.Path, m.data); err != nil {
 			return err
 		}
-		if err := EliminateBinary(b); err != nil {
+		if err := EliminateBinary(b.Path); err != nil {
 			return err
 		}
-		m.data.Binaries = append(m.data.Binaries[:idx], m.data.Binaries[idx+1:]...)
+		m.data.Binaries[idx].Eliminated = true
 	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:]...)
+		m.data.OutboundConnections[idx].Eliminated = true
 	case "directories":
 		d := m.data.Directories[idx]
-		if err := CheckDirectoryBlocked(d, m.data); err != nil {
+		if err := CheckDirectoryBlocked(d.Path, m.data); err != nil {
 			return err
 		}
-		if err := EliminateDirectory(d); err != nil {
+		if err := EliminateDirectory(d.Path); err != nil {
 			return err
 		}
-		m.data.Directories = append(m.data.Directories[:idx], m.data.Directories[idx+1:]...)
+		m.data.Directories[idx].Eliminated = true
 	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:]...)
+		m.data.Processes[idx].Eliminated = true
 	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:]...)
+		m.data.ScheduledTasks[idx].Eliminated = true
 	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:]...)
+		m.data.Services[idx].Eliminated = true
 	}
 	return nil
 }
@@ -234,6 +256,7 @@ func (m *AppModel) loadSelectedFile(path string) error {
 			return err
 		}
 		m.data = sus
+		m.loadEliminatedState()
 		return nil
 	}
 	// Try bare suspicious structure
@@ -242,9 +265,133 @@ func (m *AppModel) loadSelectedFile(path string) error {
 		return fmt.Errorf("no findings in report")
 	}
 	m.data = sus
+	m.loadEliminatedState()
 	return nil
 }

+// loadEliminatedState populates the eliminated map from the data structures
+func (m *AppModel) loadEliminatedState() {
+	m.eliminated = make(map[string]map[int]bool)
+
+	// Load eliminated autoruns
+	for i, ar := range m.data.AutoRuns {
+		if ar.Eliminated {
+			if m.eliminated["autoruns"] == nil {
+				m.eliminated["autoruns"] = make(map[int]bool)
+			}
+			m.eliminated["autoruns"][i] = true
+		}
+	}
+
+	// Load eliminated binaries
+	for i, b := range m.data.Binaries {
+		if b.Eliminated {
+			if m.eliminated["binaries"] == nil {
+				m.eliminated["binaries"] = make(map[int]bool)
+			}
+			m.eliminated["binaries"][i] = true
+		}
+	}
+
+	// Load eliminated connections
+	for i, c := range m.data.OutboundConnections {
+		if c.Eliminated {
+			if m.eliminated["connections"] == nil {
+				m.eliminated["connections"] = make(map[int]bool)
+			}
+			m.eliminated["connections"][i] = true
+		}
+	}
+
+	// Load eliminated directories
+	for i, d := range m.data.Directories {
+		if d.Eliminated {
+			if m.eliminated["directories"] == nil {
+				m.eliminated["directories"] = make(map[int]bool)
+			}
+			m.eliminated["directories"][i] = true
+		}
+	}
+
+	// Load eliminated processes
+	for i, p := range m.data.Processes {
+		if p.Eliminated {
+			if m.eliminated["processes"] == nil {
+				m.eliminated["processes"] = make(map[int]bool)
+			}
+			m.eliminated["processes"][i] = true
+		}
+	}
+
+	// Load eliminated scheduled tasks
+	for i, t := range m.data.ScheduledTasks {
+		if t.Eliminated {
+			if m.eliminated["scheduledTasks"] == nil {
+				m.eliminated["scheduledTasks"] = make(map[int]bool)
+			}
+			m.eliminated["scheduledTasks"][i] = true
+		}
+	}
+
+	// Load eliminated services
+	for i, s := range m.data.Services {
+		if s.Eliminated {
+			if m.eliminated["services"] == nil {
+				m.eliminated["services"] = make(map[int]bool)
+			}
+			m.eliminated["services"][i] = true
+		}
+	}
+}
+
+// saveDataToFile saves the current data back to the JSON file
+func (m *AppModel) saveDataToFile() error {
+	if m.filePath == "" {
+		return fmt.Errorf("no file path set")
+	}
+
+	// Read the original file to determine format
+	b, err := os.ReadFile(m.filePath)
+	if err != nil {
+		return err
+	}
+
+	// Check if it's wrapped format
+	var envelope struct {
+		Findings json.RawMessage `json:"findings"`
+	}
+	isWrapped := json.Unmarshal(b, &envelope) == nil && len(envelope.Findings) > 0
+
+	var output []byte
+	if isWrapped {
+		// Re-read the full envelope
+		var fullEnvelope map[string]interface{}
+		if err := json.Unmarshal(b, &fullEnvelope); err != nil {
+			return err
+		}
+
+		// Update the findings
+		findingsJSON, err := json.MarshalIndent(m.data, "", "  ")
+		if err != nil {
+			return err
+		}
+		fullEnvelope["findings"] = json.RawMessage(findingsJSON)
+
+		output, err = json.MarshalIndent(fullEnvelope, "", "  ")
+		if err != nil {
+			return err
+		}
+	} else {
+		// Bare format
+		output, err = json.MarshalIndent(m.data, "", "  ")
+		if err != nil {
+			return err
+		}
+	}
+
+	return os.WriteFile(m.filePath, output, 0644)
+}
+
 // RunEliminateUI starts the Bubble Tea program for elimination UI
 func RunEliminateUI() error {
 	p := tea.NewProgram(NewApp())