diff --git a/internal/tui/model.go b/internal/tui/model.go index 835e9a6..61ccfeb 100644 --- a/internal/tui/model.go +++ b/internal/tui/model.go @@ -56,24 +56,28 @@ type Model struct { height int // Network map configuration - networkMap *config.NetworkMap + networkMap *config.NetworkMap + captureMode CaptureMode + sshConfig SSHConfigModel // Capture state - captureMode CaptureMode - sshConfig SSHConfigModel - capturer *capture.Capturer - localCapturer *capture.LocalCapturer - connected bool capturing bool + connected bool + captureIface string + localCapturer *capture.LocalCapturer + capturer *capture.Capturer + captureError string packetCount int lastPackets []string - captureError string - captureIface string + + packetChan chan *sip.Packet // Channel for receiving packets from callbacks + + // Data stores + callFlowStore *sip.CallFlowStore // Call flow analysis - callFlowStore *sip.CallFlowStore - selectedFlow int - flowList list.Model + selectedFlow int + flowList list.Model // File browser for pcap import fileBrowser FileBrowserModel @@ -237,6 +241,13 @@ func NewModel() Model { } } +// waitForPacket waits for a packet on the channel +func waitForPacket(ch chan *sip.Packet) tea.Cmd { + return func() tea.Msg { + return PacketMsg{Packet: <-ch} + } +} + func createNodeInputs() []textinput.Model { inputs := make([]textinput.Model, 4) @@ -310,13 +321,20 @@ func (m Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { m.height = msg.Height case PacketMsg: - m.packetCount++ - summary := formatPacketSummary(msg.Packet, m.networkMap) - m.lastPackets = append(m.lastPackets, summary) - if len(m.lastPackets) > 50 { - m.lastPackets = m.lastPackets[1:] + if msg.Packet != nil { + m.packetCount++ + summary := formatPacketSummary(msg.Packet, m.networkMap) + m.lastPackets = append(m.lastPackets, summary) + if len(m.lastPackets) > 50 { + m.lastPackets = m.lastPackets[1:] + } + m.callFlowStore.AddPacket(msg.Packet) + } + + // Continue waiting if we are still capturing + if m.capturing { + return m, waitForPacket(m.packetChan) } - m.callFlowStore.AddPacket(msg.Packet) case ErrorMsg: m.captureError = msg.Error.Error() @@ -560,9 +578,14 @@ func (m *Model) startLocalCapture() tea.Cmd { m.lastPackets = m.lastPackets[:0] m.captureMode = CaptureModeLocal + // Create a buffered channel for packets + m.packetChan = make(chan *sip.Packet, 100) + m.localCapturer = capture.NewLocalCapturer() m.localCapturer.OnPacket = func(p *sip.Packet) { - // Note: In real implementation, use channel + tea.Cmd + if m.capturing { + m.packetChan <- p + } } m.localCapturer.OnError = func(err error) { m.captureError = err.Error() @@ -576,9 +599,10 @@ func (m *Model) startLocalCapture() tea.Cmd { if err := m.localCapturer.Start(iface, 5060); err != nil { m.captureError = err.Error() m.capturing = false + return nil } - return nil + return waitForPacket(m.packetChan) } func (m *Model) startSSHCapture() tea.Cmd { @@ -587,8 +611,13 @@ func (m *Model) startSSHCapture() tea.Cmd { m.packetCount = 0 m.lastPackets = m.lastPackets[:0] + // Create a buffered channel for packets + m.packetChan = make(chan *sip.Packet, 100) + m.capturer.OnPacket = func(p *sip.Packet) { - // Note: In real implementation, use channel + tea.Cmd + if m.capturing { + m.packetChan <- p + } } m.capturer.OnError = func(err error) { m.captureError = err.Error() @@ -597,9 +626,10 @@ func (m *Model) startSSHCapture() tea.Cmd { if err := m.capturer.Start("any", 5060); err != nil { m.captureError = err.Error() m.capturing = false + return nil } - return nil + return waitForPacket(m.packetChan) } func (m *Model) stopCapture() { @@ -609,7 +639,16 @@ func (m *Model) stopCapture() { if m.capturer != nil { m.capturer.Stop() } + + wasCapturing := m.capturing m.capturing = false + + // Unblock any waiting waitForPacket command + if wasCapturing && m.packetChan != nil { + go func() { + m.packetChan <- nil + }() + } } func (m *Model) disconnect() {