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}