mirror of
https://github.com/usestrix/strix.git
synced 2026-08-16 09:26:39 +02:00
468 lines
14 KiB
Go
468 lines
14 KiB
Go
package app
|
|
|
|
import (
|
|
"encoding/json"
|
|
"strings"
|
|
"time"
|
|
|
|
tea "github.com/charmbracelet/bubbletea"
|
|
"github.com/usestrix/strix/tui/internal/protocol"
|
|
"github.com/usestrix/strix/tui/internal/render"
|
|
)
|
|
|
|
func (m *Model) handleEnvelope(envelope protocol.Envelope) tea.Cmd {
|
|
switch envelope.Type {
|
|
case "state":
|
|
var update protocol.StateUpdate
|
|
if err := json.Unmarshal(envelope.Payload, &update); err != nil {
|
|
m.errorText = err.Error()
|
|
return nil
|
|
}
|
|
if update.Revision <= m.stateRevision {
|
|
return nil
|
|
}
|
|
selectedAgentID := ""
|
|
if m.selectedAgent >= 0 && m.selectedAgent < len(m.snapshot.Agents) {
|
|
selectedAgentID = m.snapshot.Agents[m.selectedAgent].ID
|
|
}
|
|
update.State.Events = m.snapshot.Events
|
|
update.State.Vulnerabilities = m.snapshot.Vulnerabilities
|
|
update.State.Agents = m.snapshot.Agents
|
|
m.consumeMessages(update.State.Messages, update.State.SetupMode)
|
|
m.snapshot = update.State
|
|
m.stateRevision = update.Revision
|
|
if m.snapshot.Error != nil {
|
|
m.errorText = *m.snapshot.Error
|
|
}
|
|
if m.snapshot.SetupMode {
|
|
// The start screen is its own landing page; never sit on the
|
|
// splash before it.
|
|
m.showSplash = false
|
|
m.input.Placeholder = setupPlaceholder
|
|
} else {
|
|
m.input.Placeholder = chatPlaceholder
|
|
}
|
|
m.selectedAgent = selectedAgentIndex(m.snapshot.Agents, selectedAgentID)
|
|
m.selectedVuln = min(m.selectedVuln, max(0, len(m.snapshot.Vulnerabilities)-1))
|
|
if m.modal == modalStop && !m.selectedAgentCanStop() {
|
|
m.closeModal()
|
|
}
|
|
m.syncMountPrompt()
|
|
m.ensureAgentVisible()
|
|
m.ensureVulnerabilityVisible()
|
|
m.ready = true
|
|
// resize (not just refresh): status-row visibility changes the chat height.
|
|
m.resizeViewport()
|
|
m.resizeVulnerabilityViewport()
|
|
case "collection_bootstrap":
|
|
return m.handleCollectionBootstrap(envelope.Payload)
|
|
case "collection_delta":
|
|
return m.handleCollectionDelta(envelope.Payload)
|
|
case "command_result":
|
|
if m.client == nil {
|
|
return nil
|
|
}
|
|
expectedCommand, pending := m.client.ExpectedCommand(envelope.RequestID)
|
|
if !pending {
|
|
return nil
|
|
}
|
|
var result protocol.CommandResult
|
|
if err := json.Unmarshal(envelope.Payload, &result); err != nil {
|
|
m.errorText = err.Error()
|
|
return nil
|
|
}
|
|
if result.Command != expectedCommand || !m.client.Resolve(envelope.RequestID, result.Command) {
|
|
return nil
|
|
}
|
|
if !result.OK {
|
|
if result.Command == "collection.resync" {
|
|
if collection := m.resyncRequests[envelope.RequestID]; collection != "" {
|
|
m.resyncRequested[collection] = false
|
|
delete(m.resyncRequests, envelope.RequestID)
|
|
}
|
|
}
|
|
message := "Command failed"
|
|
if result.Error != nil && strings.TrimSpace(result.Error.Message) != "" {
|
|
message = result.Error.Message
|
|
}
|
|
// Setup-mode errors live in the scrollback (red), like Python; during
|
|
// a scan they surface on the status line.
|
|
if m.snapshot.SetupMode {
|
|
m.setupMsg(message, render.Col(red))
|
|
} else {
|
|
m.errorText = message
|
|
}
|
|
return nil
|
|
}
|
|
if m.snapshot.ScanStarted && !m.snapshot.SetupMode && strings.HasPrefix(result.Command, "setup.") {
|
|
return nil
|
|
}
|
|
m.errorText = ""
|
|
switch result.Command {
|
|
case "viewer.open":
|
|
var data struct {
|
|
Status string `json:"status"`
|
|
URL *string `json:"url"`
|
|
}
|
|
_ = json.Unmarshal(result.Result, &data)
|
|
m.snapshot.ViewerStatus = data.Status
|
|
m.snapshot.ViewerURL = data.URL
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (m *Model) consumeMessages(messages []protocol.Message, setupMode bool) {
|
|
if m.seenMessages == nil {
|
|
m.seenMessages = map[string]bool{}
|
|
}
|
|
for _, message := range messages {
|
|
key := message.ID
|
|
if key == "" {
|
|
key = message.Level + "\x00" + message.Text
|
|
}
|
|
if m.seenMessages[key] {
|
|
continue
|
|
}
|
|
m.seenMessages[key] = true
|
|
if !setupMode || strings.TrimSpace(message.Text) == "" {
|
|
continue
|
|
}
|
|
style := render.Dim()
|
|
switch message.Level {
|
|
case "error":
|
|
style = render.Col(red)
|
|
case "warning":
|
|
style = render.Col(amber)
|
|
}
|
|
m.setupMsg(message.Text, style)
|
|
}
|
|
}
|
|
|
|
func validCollection(name string) bool {
|
|
return name == "agents" || name == "events" || name == "vulnerabilities"
|
|
}
|
|
|
|
func (m *Model) collectionMismatch(name string) tea.Cmd {
|
|
delete(m.collectionAssemblies, name)
|
|
if !validCollection(name) || m.resyncRequested[name] || m.client == nil {
|
|
return nil
|
|
}
|
|
m.resyncRequested[name] = true
|
|
return send(m.client, "collection.resync", map[string]any{"collection": name})
|
|
}
|
|
|
|
func (m *Model) clearCollectionResync(name string) {
|
|
m.resyncRequested[name] = false
|
|
for requestID, collection := range m.resyncRequests {
|
|
if collection == name {
|
|
delete(m.resyncRequests, requestID)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (m *Model) handleCollectionBootstrap(payload json.RawMessage) tea.Cmd {
|
|
var chunk protocol.CollectionBootstrap
|
|
if err := json.Unmarshal(payload, &chunk); err != nil {
|
|
m.errorText = err.Error()
|
|
return nil
|
|
}
|
|
if !validCollection(chunk.Collection) {
|
|
m.errorText = "Unknown collection: " + chunk.Collection
|
|
return nil
|
|
}
|
|
if chunk.Cursor == 0 {
|
|
m.resyncRequested[chunk.Collection] = false
|
|
}
|
|
if chunk.Cursor == 0 {
|
|
if chunk.Revision <= m.collectionRevisions[chunk.Collection] {
|
|
return nil
|
|
}
|
|
m.collectionAssemblies[chunk.Collection] = &collectionAssembly{
|
|
kind: "bootstrap", revision: chunk.Revision, ids: map[string]bool{},
|
|
}
|
|
}
|
|
assembly := m.collectionAssemblies[chunk.Collection]
|
|
if assembly == nil || assembly.kind != "bootstrap" || assembly.revision != chunk.Revision || assembly.cursor != chunk.Cursor {
|
|
return m.collectionMismatch(chunk.Collection)
|
|
}
|
|
if chunk.NextCursor != chunk.Cursor+len(chunk.Items) {
|
|
return m.collectionMismatch(chunk.Collection)
|
|
}
|
|
for _, raw := range chunk.Items {
|
|
if chunk.Collection == "agents" {
|
|
var agent protocol.Agent
|
|
if err := json.Unmarshal(raw, &agent); err != nil || agent.ID == "" {
|
|
return m.collectionMismatch(chunk.Collection)
|
|
}
|
|
if assembly.ids[agent.ID] {
|
|
return m.collectionMismatch(chunk.Collection)
|
|
}
|
|
assembly.ids[agent.ID] = true
|
|
assembly.agents = append(assembly.agents, agent)
|
|
} else if chunk.Collection == "events" {
|
|
var event protocol.Event
|
|
if err := json.Unmarshal(raw, &event); err != nil || event.ID == "" {
|
|
return m.collectionMismatch(chunk.Collection)
|
|
}
|
|
if assembly.ids[event.ID] {
|
|
return m.collectionMismatch(chunk.Collection)
|
|
}
|
|
assembly.ids[event.ID] = true
|
|
assembly.events = append(assembly.events, event)
|
|
} else {
|
|
var finding map[string]any
|
|
if err := json.Unmarshal(raw, &finding); err != nil || collectionItemID(finding) == "" {
|
|
return m.collectionMismatch(chunk.Collection)
|
|
}
|
|
id := collectionItemID(finding)
|
|
if assembly.ids[id] {
|
|
return m.collectionMismatch(chunk.Collection)
|
|
}
|
|
assembly.ids[id] = true
|
|
assembly.findings = append(assembly.findings, finding)
|
|
}
|
|
}
|
|
assembly.cursor = chunk.NextCursor
|
|
if !chunk.Done {
|
|
return nil
|
|
}
|
|
if chunk.Collection == "agents" {
|
|
selectedAgentID := m.selectedAgentID()
|
|
m.snapshot.Agents = assembly.agents
|
|
m.selectedAgent = selectedAgentIndex(m.snapshot.Agents, selectedAgentID)
|
|
} else if chunk.Collection == "events" {
|
|
m.snapshot.Events = assembly.events
|
|
} else {
|
|
m.snapshot.Vulnerabilities = assembly.findings
|
|
}
|
|
m.collectionRevisions[chunk.Collection] = chunk.Revision
|
|
delete(m.collectionAssemblies, chunk.Collection)
|
|
m.clearCollectionResync(chunk.Collection)
|
|
return m.refreshAfterCollection(chunk.Collection)
|
|
}
|
|
|
|
func (m *Model) handleCollectionDelta(payload json.RawMessage) tea.Cmd {
|
|
var chunk protocol.CollectionDelta
|
|
if err := json.Unmarshal(payload, &chunk); err != nil {
|
|
m.errorText = err.Error()
|
|
return nil
|
|
}
|
|
if !validCollection(chunk.Collection) {
|
|
m.errorText = "Unknown collection: " + chunk.Collection
|
|
return nil
|
|
}
|
|
if chunk.Cursor == 0 {
|
|
if chunk.BaseRevision != m.collectionRevisions[chunk.Collection] || chunk.Revision <= chunk.BaseRevision {
|
|
return m.collectionMismatch(chunk.Collection)
|
|
}
|
|
m.collectionAssemblies[chunk.Collection] = &collectionAssembly{
|
|
kind: "delta", revision: chunk.Revision, baseRevision: chunk.BaseRevision,
|
|
}
|
|
}
|
|
assembly := m.collectionAssemblies[chunk.Collection]
|
|
if assembly == nil || assembly.kind != "delta" || assembly.revision != chunk.Revision ||
|
|
assembly.baseRevision != chunk.BaseRevision || assembly.cursor != chunk.Cursor {
|
|
return m.collectionMismatch(chunk.Collection)
|
|
}
|
|
if chunk.NextCursor != chunk.Cursor+len(chunk.Operations) {
|
|
return m.collectionMismatch(chunk.Collection)
|
|
}
|
|
assembly.operations = append(assembly.operations, chunk.Operations...)
|
|
assembly.cursor = chunk.NextCursor
|
|
if !chunk.Done {
|
|
return nil
|
|
}
|
|
if !m.applyCollectionOperations(chunk.Collection, assembly.operations) {
|
|
return m.collectionMismatch(chunk.Collection)
|
|
}
|
|
m.collectionRevisions[chunk.Collection] = chunk.Revision
|
|
delete(m.collectionAssemblies, chunk.Collection)
|
|
m.clearCollectionResync(chunk.Collection)
|
|
return m.refreshAfterCollection(chunk.Collection)
|
|
}
|
|
|
|
func (m *Model) applyCollectionOperations(name string, operations []protocol.CollectionOperation) bool {
|
|
seen := make(map[string]bool, len(operations))
|
|
if name == "agents" {
|
|
selectedAgentID := m.selectedAgentID()
|
|
values := append([]protocol.Agent(nil), m.snapshot.Agents...)
|
|
positions := make(map[string]int, len(values))
|
|
for index, agent := range values {
|
|
positions[agent.ID] = index
|
|
}
|
|
for _, operation := range operations {
|
|
if operation.Op == "delete" {
|
|
if operation.ID == "" || seen[operation.ID] {
|
|
return false
|
|
}
|
|
seen[operation.ID] = true
|
|
index, exists := positions[operation.ID]
|
|
if !exists {
|
|
return false
|
|
}
|
|
values = append(values[:index], values[index+1:]...)
|
|
positions = make(map[string]int, len(values))
|
|
for position, value := range values {
|
|
positions[value.ID] = position
|
|
}
|
|
continue
|
|
}
|
|
if operation.Op != "upsert" {
|
|
return false
|
|
}
|
|
var agent protocol.Agent
|
|
if err := json.Unmarshal(operation.Item, &agent); err != nil || agent.ID == "" || seen[agent.ID] {
|
|
return false
|
|
}
|
|
seen[agent.ID] = true
|
|
if index, exists := positions[agent.ID]; exists {
|
|
values[index] = agent
|
|
} else {
|
|
positions[agent.ID] = len(values)
|
|
values = append(values, agent)
|
|
}
|
|
}
|
|
m.snapshot.Agents = values
|
|
m.selectedAgent = selectedAgentIndex(values, selectedAgentID)
|
|
return true
|
|
}
|
|
if name == "events" {
|
|
values := append([]protocol.Event(nil), m.snapshot.Events...)
|
|
positions := make(map[string]int, len(values))
|
|
for index, event := range values {
|
|
positions[event.ID] = index
|
|
}
|
|
for _, operation := range operations {
|
|
if operation.Op == "delete" {
|
|
if operation.ID == "" || seen[operation.ID] {
|
|
return false
|
|
}
|
|
seen[operation.ID] = true
|
|
index, exists := positions[operation.ID]
|
|
if !exists {
|
|
return false
|
|
}
|
|
values = append(values[:index], values[index+1:]...)
|
|
positions = make(map[string]int, len(values))
|
|
for position, value := range values {
|
|
positions[value.ID] = position
|
|
}
|
|
continue
|
|
}
|
|
if operation.Op != "upsert" {
|
|
return false
|
|
}
|
|
var event protocol.Event
|
|
if err := json.Unmarshal(operation.Item, &event); err != nil || event.ID == "" || event.Version < 0 || seen[event.ID] {
|
|
return false
|
|
}
|
|
seen[event.ID] = true
|
|
if index, exists := positions[event.ID]; exists {
|
|
current := values[index]
|
|
if event.Version <= current.Version {
|
|
return false
|
|
}
|
|
values[index] = event
|
|
} else {
|
|
positions[event.ID] = len(values)
|
|
values = append(values, event)
|
|
}
|
|
}
|
|
m.snapshot.Events = values
|
|
return true
|
|
}
|
|
|
|
values := append([]map[string]any(nil), m.snapshot.Vulnerabilities...)
|
|
positions := make(map[string]int, len(values))
|
|
for index, finding := range values {
|
|
positions[collectionItemID(finding)] = index
|
|
}
|
|
for _, operation := range operations {
|
|
if operation.Op == "delete" {
|
|
if operation.ID == "" || seen[operation.ID] {
|
|
return false
|
|
}
|
|
seen[operation.ID] = true
|
|
index, exists := positions[operation.ID]
|
|
if !exists {
|
|
return false
|
|
}
|
|
values = append(values[:index], values[index+1:]...)
|
|
positions = make(map[string]int, len(values))
|
|
for position, value := range values {
|
|
positions[collectionItemID(value)] = position
|
|
}
|
|
continue
|
|
}
|
|
if operation.Op != "upsert" {
|
|
return false
|
|
}
|
|
var finding map[string]any
|
|
if err := json.Unmarshal(operation.Item, &finding); err != nil {
|
|
return false
|
|
}
|
|
id := collectionItemID(finding)
|
|
if id == "" {
|
|
return false
|
|
}
|
|
if seen[id] {
|
|
return false
|
|
}
|
|
seen[id] = true
|
|
if index, exists := positions[id]; exists {
|
|
values[index] = finding
|
|
} else {
|
|
positions[id] = len(values)
|
|
values = append(values, finding)
|
|
}
|
|
}
|
|
m.snapshot.Vulnerabilities = values
|
|
return true
|
|
}
|
|
|
|
func collectionItemID(item map[string]any) string {
|
|
id, _ := item["id"].(string)
|
|
return id
|
|
}
|
|
|
|
func (m *Model) refreshAfterCollection(name string) tea.Cmd {
|
|
if name == "agents" {
|
|
m.ensureAgentVisible()
|
|
m.refreshViewport()
|
|
return m.notifyBudgetPause()
|
|
}
|
|
if name == "events" {
|
|
m.refreshViewport()
|
|
return nil
|
|
}
|
|
m.selectedVuln = min(m.selectedVuln, max(0, len(m.snapshot.Vulnerabilities)-1))
|
|
m.ensureVulnerabilityVisible()
|
|
m.resizeVulnerabilityViewport()
|
|
return nil
|
|
}
|
|
|
|
// notifyBudgetPause ports _notify_budget_pause: a one-shot warning toast when
|
|
// any agent hits the budget limit, re-armed once no agent is paused.
|
|
func (m *Model) notifyBudgetPause() tea.Cmd {
|
|
paused := false
|
|
for _, agent := range m.snapshot.Agents {
|
|
if agent.Status == "budget_paused" {
|
|
paused = true
|
|
break
|
|
}
|
|
}
|
|
if paused && !m.budgetPauseNotified {
|
|
m.budgetPauseNotified = true
|
|
return m.showToastFor(
|
|
"Budget limit reached — agents paused. Send a message to continue "+
|
|
"(this extends the budget), or ctrl-q to quit.",
|
|
15*time.Second,
|
|
)
|
|
}
|
|
if !paused {
|
|
m.budgetPauseNotified = false
|
|
}
|
|
return nil
|
|
}
|