main
6d86999 ยท 6 months ago 26 commits
  1package rules
  2
  3import (
  4	"log"
  5	"github.com/The-Robin-Hood/ufWall/internal/keys"
  6	"github.com/The-Robin-Hood/ufWall/internal/ufw"
  7	"github.com/The-Robin-Hood/ufWall/internal/ui"
  8
  9	"github.com/charmbracelet/bubbles/key"
 10	tea "github.com/charmbracelet/bubbletea"
 11)
 12
 13const (
 14	ActionToggle   = "toggle"
 15	ActionMoveUp   = "move_up"
 16	ActionMoveDown = "move_down"
 17	ActionAdd      = "add"
 18)
 19
 20func (m Model) Update(msg tea.Msg, data RulesData) (Model, tea.Cmd) {
 21	rules := m.getActiveRules(data)
 22
 23	if m.addWizard != nil {
 24		return m.updateAddWizard(msg)
 25	}
 26
 27	if m.showDeleteConfirm {
 28		switch msg := msg.(type) {
 29		case tea.KeyMsg:
 30			switch {
 31			case key.Matches(msg, keys.Bindings.Execute):
 32				if m.deleteRule != nil {
 33					log.Printf("Deleting rule #%d (confirmed)", m.deleteRule.Num)
 34					ufw.DeleteRule(m.deleteRule.Num)
 35				}
 36				m.showDeleteConfirm = false
 37				m.deleteRule = nil
 38				return m, keys.Refresh()
 39			case key.Matches(msg, keys.Bindings.Quit):
 40				m.showDeleteConfirm = false
 41				m.deleteRule = nil
 42				return m, nil
 43			}
 44		}
 45		return m, nil
 46	}
 47
 48	if m.showDetails {
 49		switch msg.(type) {
 50		case tea.KeyMsg:
 51			m.showDetails = false
 52			m.detailRule = nil
 53			return m, nil
 54		}
 55		return m, nil
 56	}
 57
 58	if m.menu != nil {
 59		if quit := m.menu.Update(msg); quit {
 60			if m.menuContext != nil && m.menuContext.PendingSubmenu {
 61				return m.handlePendingSubmenu()
 62			}
 63			m.menu = nil
 64			m.menuContext = nil
 65			return m, keys.Refresh()
 66		}
 67		return m, nil
 68	}
 69
 70	switch msg := msg.(type) {
 71	case tea.KeyMsg:
 72		switch {
 73		case key.Matches(msg, keys.Bindings.SwitchTable):
 74			if m.activeTable == IPv4Table {
 75				m.activeTable = IPv6Table
 76			} else {
 77				m.activeTable = IPv4Table
 78			}
 79			return m, nil
 80
 81		case key.Matches(msg, keys.Bindings.CursorUp):
 82			m.moveCursorUp()
 83			return m, nil
 84
 85		case key.Matches(msg, keys.Bindings.CursorDown):
 86			m.moveCursorDown(rules)
 87			return m, nil
 88
 89		case key.Matches(msg, keys.Bindings.Execute):
 90			if len(rules) > 0 {
 91				return m.openMainMenu(data)
 92			}
 93			m.addWizard = NewAddWizard()
 94			return m, nil
 95
 96		case key.Matches(msg, keys.Bindings.AddRule):
 97			m.addWizard = NewAddWizard()
 98			return m, nil
 99
100		case key.Matches(msg, keys.Bindings.Info):
101			cursorLine := m.getCurrentCursor()
102			if len(rules) > 0 && cursorLine < len(rules) {
103				ruleCopy := rules[cursorLine]
104				m.showDetails = true
105				m.detailRule = &ruleCopy
106				return m, nil
107			}
108
109		case key.Matches(msg, keys.Bindings.Delete):
110			cursorLine := m.getCurrentCursor()
111			if len(rules) > 0 && cursorLine < len(rules) {
112				ruleCopy := rules[cursorLine]
113				m.showDeleteConfirm = true
114				m.deleteRule = &ruleCopy
115				return m, nil
116			}
117		}
118	}
119	return m, nil
120}
121
122func (m Model) updateAddWizard(msg tea.Msg) (Model, tea.Cmd) {
123	w := m.addWizard
124
125	switch msg := msg.(type) {
126	case tea.KeyMsg:
127		switch {
128		case key.Matches(msg, keys.Bindings.Quit):
129			if w.InputMode {
130				w.InputMode = false
131				w.Input = ""
132				return m, nil
133			}
134			m.addWizard = nil
135			return m, nil
136		}
137
138		if w.InputMode {
139			switch {
140			case key.Matches(msg, keys.Bindings.Execute):
141				return m.confirmWizardInput()
142			case key.Matches(msg, keys.Bindings.Back):
143				if len(w.Input) > 0 {
144					w.Input = w.Input[:len(w.Input)-1]
145				}
146			default:
147				if len(msg.String()) == 1 {
148					w.Input += msg.String()
149				}
150			}
151			return m, nil
152		}
153
154		switch {
155		case key.Matches(msg, keys.Bindings.CursorUp):
156			if w.Cursor > 0 {
157				w.Cursor--
158			}
159		case key.Matches(msg, keys.Bindings.CursorDown):
160			if w.Cursor < len(w.Options)-1 {
161				w.Cursor++
162			}
163		case key.Matches(msg, keys.Bindings.Execute):
164			return m.selectWizardOption()
165		case key.Matches(msg, keys.Bindings.CustomInput):
166			if w.Step == StepPort || w.Step == StepSource || w.Step == StepDestination {
167				w.InputMode = true
168				w.Input = ""
169			}
170		case key.Matches(msg, keys.Bindings.Back):
171			return m.wizardPrevStep()
172		}
173	}
174
175	return m, nil
176}
177
178func (m Model) selectWizardOption() (Model, tea.Cmd) {
179	w := m.addWizard
180	if w.Cursor >= len(w.Options) {
181		return m, nil
182	}
183
184	selected := w.Options[w.Cursor]
185
186	switch w.Step {
187	case StepAction:
188		w.Params.Action = selected
189		w.Step = StepDirection
190		w.Options = append([]string{"Both (in & out)"}, ufw.Directions...)
191		w.Cursor = 0
192
193	case StepDirection:
194		if selected == "Both (in & out)" {
195			w.Params.Direction = ""
196		} else {
197			w.Params.Direction = selected
198		}
199		w.Step = StepProtocol
200		w.Options = ufw.Protocols
201		w.Cursor = 0
202
203	case StepProtocol:
204		w.Params.Protocol = selected
205		w.Step = StepPort
206		w.Options = []string{"Any (no port filter)"}
207		for _, p := range ufw.CommonPorts {
208			w.Options = append(w.Options, p.Port+" ("+p.Name+")")
209		}
210		w.Cursor = 0
211
212	case StepPort:
213		if selected == "Any (no port filter)" {
214			w.Params.Port = ""
215		} else {
216			for i, c := range selected {
217				if c == ' ' {
218					w.Params.Port = selected[:i]
219					break
220				}
221			}
222		}
223		w.Step = StepSource
224		w.Options = []string{"Any", "Custom..."}
225		w.Cursor = 0
226
227	case StepSource:
228		switch selected {
229		case "Any":
230			w.Params.FromAddr = ""
231		case "Custom...":
232			w.InputMode = true
233			w.Input = ""
234			return m, nil
235		}
236		w.Step = StepDestination
237		w.Options = []string{"Any", "Custom..."}
238		w.Cursor = 0
239
240	case StepDestination:
241		switch selected {
242		case "Any":
243			w.Params.ToAddr = ""
244		case "Custom...":
245			w.InputMode = true
246			w.Input = ""
247			return m, nil
248		}
249		w.Step = StepInterface
250		w.Options = append([]string{"All interfaces"}, w.Interfaces...)
251		w.Cursor = 0
252
253	case StepInterface:
254		if selected == "All interfaces" {
255			w.Params.Interface = ""
256		} else {
257			w.Params.Interface = selected
258		}
259		w.Step = StepConfirm
260		w.Options = []string{"Confirm and Add Rule", "Cancel"}
261		w.Cursor = 0
262
263	case StepConfirm:
264		if selected == "Confirm and Add Rule" {
265			// Execute the command
266			log.Printf("Adding rule: %+v", w.Params)
267			_, stderr, err := ufw.AddNewRule(w.Params)
268			if err != nil {
269				log.Printf("Error adding rule: %s", stderr)
270				w.Error = stderr
271				return m, nil
272			}
273			m.addWizard = nil
274			return m, keys.Refresh()
275		} else {
276			// Cancel
277			m.addWizard = nil
278			return m, nil
279		}
280	}
281
282	return m, nil
283}
284
285func (m Model) confirmWizardInput() (Model, tea.Cmd) {
286	w := m.addWizard
287	input := w.Input
288	w.InputMode = false
289	w.Input = ""
290
291	switch w.Step {
292	case StepPort:
293		w.Params.Port = input
294		w.Step = StepSource
295		w.Options = []string{"Any", "Custom..."}
296		w.Cursor = 0
297
298	case StepSource:
299		w.Params.FromAddr = input
300		w.Step = StepDestination
301		w.Options = []string{"Any", "Custom..."}
302		w.Cursor = 0
303
304	case StepDestination:
305		w.Params.ToAddr = input
306		w.Step = StepInterface
307		w.Options = append([]string{"All interfaces"}, w.Interfaces...)
308		w.Cursor = 0
309	}
310
311	return m, nil
312}
313
314func (m Model) wizardPrevStep() (Model, tea.Cmd) {
315	w := m.addWizard
316
317	switch w.Step {
318	case StepAction:
319		m.addWizard = nil
320		return m, nil
321
322	case StepDirection:
323		w.Step = StepAction
324		w.Options = ufw.Actions
325		w.Cursor = 0
326
327	case StepProtocol:
328		w.Step = StepDirection
329		w.Options = append([]string{"Both (in & out)"}, ufw.Directions...)
330		w.Cursor = 0
331
332	case StepPort:
333		w.Step = StepProtocol
334		w.Options = ufw.Protocols
335		w.Cursor = 0
336
337	case StepSource:
338		w.Step = StepPort
339		w.Options = []string{"Any (no port filter)"}
340		for _, p := range ufw.CommonPorts {
341			w.Options = append(w.Options, p.Port+" ("+p.Name+")")
342		}
343		w.Cursor = 0
344
345	case StepDestination:
346		w.Step = StepSource
347		w.Options = []string{"Any", "Custom..."}
348		w.Cursor = 0
349
350	case StepInterface:
351		w.Step = StepDestination
352		w.Options = []string{"Any", "Custom..."}
353		w.Cursor = 0
354
355	case StepConfirm:
356		w.Step = StepInterface
357		w.Options = append([]string{"All interfaces"}, w.Interfaces...)
358		w.Cursor = 0
359	}
360
361	return m, nil
362}
363
364func (m Model) getActiveRules(data RulesData) []ufw.Rule {
365	if m.activeTable == IPv6Table {
366		return data.IPv6
367	}
368	return data.IPv4
369}
370
371func (m Model) getCurrentCursor() int {
372	if m.activeTable == IPv6Table {
373		return m.ipv6CursorLine
374	}
375	return m.ipv4CursorLine
376}
377
378func (m *Model) moveCursorUp() {
379	if m.activeTable == IPv6Table {
380		if m.ipv6CursorLine > 0 {
381			m.ipv6CursorLine--
382			// Scroll up if cursor goes above visible area
383			if m.ipv6CursorLine < m.ipv6ScrollOffset {
384				m.ipv6ScrollOffset = m.ipv6CursorLine
385			}
386		}
387	} else {
388		if m.ipv4CursorLine > 0 {
389			m.ipv4CursorLine--
390			// Scroll up if cursor goes above visible area
391			if m.ipv4CursorLine < m.ipv4ScrollOffset {
392				m.ipv4ScrollOffset = m.ipv4CursorLine
393			}
394		}
395	}
396}
397
398func (m *Model) moveCursorDown(rules []ufw.Rule) {
399	if m.activeTable == IPv6Table {
400		if m.ipv6CursorLine < len(rules)-1 {
401			m.ipv6CursorLine++
402			// Scroll down if cursor goes below visible area
403			if m.ipv6CursorLine >= m.ipv6ScrollOffset+MaxVisibleRules {
404				m.ipv6ScrollOffset = m.ipv6CursorLine - MaxVisibleRules + 1
405			}
406		}
407	} else {
408		if m.ipv4CursorLine < len(rules)-1 {
409			m.ipv4CursorLine++
410			// Scroll down if cursor goes below visible area
411			if m.ipv4CursorLine >= m.ipv4ScrollOffset+MaxVisibleRules {
412				m.ipv4ScrollOffset = m.ipv4CursorLine - MaxVisibleRules + 1
413			}
414		}
415	}
416}
417
418func (m Model) handlePendingSubmenu() (Model, tea.Cmd) {
419	if m.menuContext == nil {
420		m.menu = nil
421		return m, keys.Refresh()
422	}
423
424	log.Printf("handlePendingSubmenu: Action=%s, PendingSubmenu=%v", m.menuContext.Action, m.menuContext.PendingSubmenu)
425	m.menuContext.PendingSubmenu = false
426
427	switch m.menuContext.Action {
428	case ActionToggle:
429		return m.openToggleActionMenu()
430
431	case ActionMoveUp:
432		rule := m.menuContext.SelectedRule
433		if rule != nil {
434			log.Printf("Moving rule #%d up (IPv6=%v)", rule.Num, rule.IPv6)
435			err := ufw.MoveRule(*rule, -1, rule.Action)
436			if err != nil {
437				log.Printf("Error moving rule up: %v", err)
438			}
439		}
440		m.menu = nil
441		m.menuContext = nil
442		return m, keys.Refresh()
443
444	case ActionMoveDown:
445		rule := m.menuContext.SelectedRule
446		if rule != nil {
447			log.Printf("Moving rule #%d down (IPv6=%v)", rule.Num, rule.IPv6)
448			err := ufw.MoveRule(*rule, 1, rule.Action)
449			if err != nil {
450				log.Printf("Error moving rule down: %v", err)
451			}
452		}
453		m.menu = nil
454		m.menuContext = nil
455		return m, keys.Refresh()
456
457	case ActionAdd:
458		m.menu = nil
459		m.menuContext = nil
460		m.addWizard = NewAddWizard()
461		return m, nil
462	}
463
464	m.menu = nil
465	m.menuContext = nil
466	return m, keys.Refresh()
467}
468
469func (m Model) openMainMenu(data RulesData) (Model, tea.Cmd) {
470	rules := m.getActiveRules(data)
471	cursorLine := m.getCurrentCursor()
472
473	if len(rules) == 0 || cursorLine >= len(rules) {
474		return m, nil
475	}
476
477	selectedRule := rules[cursorLine]
478	ruleCopy := selectedRule // Make a copy to store in context
479
480	menuLabels := []string{
481		"Toggle Action",
482		"Move Up",
483		"Move Down",
484		"Add New Rule",
485	}
486
487	m.menuContext = &MenuContext{
488		SelectedRule: &ruleCopy,
489		TotalRules:   len(rules),
490	}
491
492	options := ui.MakeMenuItems(menuLabels, func(label string) tea.Cmd {
493		switch label {
494		case "Toggle Action":
495			m.menuContext.Action = ActionToggle
496		case "Move Up":
497			m.menuContext.Action = ActionMoveUp
498		case "Move Down":
499			m.menuContext.Action = ActionMoveDown
500		case "Add New Rule":
501			m.menuContext.Action = ActionAdd
502		}
503		m.menuContext.PendingSubmenu = true
504		return nil
505	})
506
507	menu := ui.NewMenu(options, m.styles)
508	m.menu = &menu
509	return m, nil
510}
511
512func (m Model) openToggleActionMenu() (Model, tea.Cmd) {
513	rule := m.menuContext.SelectedRule
514
515	options := ui.MakeMenuItems(ufw.Actions, func(newAction string) tea.Cmd {
516		if rule != nil && newAction != rule.Action {
517			log.Printf("Toggling rule #%d from %s to %s (IPv6=%v)", rule.Num, rule.Action, newAction, rule.IPv6)
518
519			err := ufw.MoveRule(*rule, 0, newAction)
520			if err != nil {
521				log.Printf("Error toggling rule: %v", err)
522			}
523		}
524		m.menuContext = nil
525		return nil
526	})
527
528	menu := ui.NewMenu(options, m.styles)
529	m.menu = &menu
530	return m, nil
531}