Skip to content

Commit 177c49a

Browse files
pelthopeltho
committed
Add Edit rule feature
1 parent 300f494 commit 177c49a

File tree

1 file changed

+136
-52
lines changed

1 file changed

+136
-52
lines changed

‎main.go

+136-52
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55
"fmt"
66
"log"
77
"os/exec"
8+
"regexp"
89
"strings"
910

1011
"github.com/gdamore/tcell/v2"
@@ -105,7 +106,73 @@ func (t *Tui) CreateTable(rows []string) {
105106
t.help.Clear()
106107
t.app.SetFocus(t.menu)
107108
}
108-
}).SetSelectedFunc(func(row int, column int) {
109+
})
110+
}
111+
112+
func (t *Tui) ReloadTable() {
113+
t.table.Clear()
114+
data, _ := t.LoadTableData()
115+
t.CreateTable(data)
116+
}
117+
118+
func (t *Tui) CreateModal(text string, confirm func(), cancel func(), finally func()) {
119+
modal := tview.NewModal()
120+
t.pages.AddPage("modal", modal.SetText(text).AddButtons([]string{"Confirm", "Cancel"}).SetDoneFunc(func(i int, label string) {
121+
if label == "Confirm" {
122+
confirm()
123+
t.ReloadTable()
124+
} else {
125+
cancel()
126+
}
127+
modal.ClearButtons()
128+
finally()
129+
}), true, true)
130+
}
131+
132+
func (t *Tui) CreateRule(position ...int) {
133+
to := t.form.GetFormItem(0).(*tview.InputField).GetText()
134+
port := t.form.GetFormItem(1).(*tview.InputField).GetText()
135+
_, proto := t.form.GetFormItem(2).(*tview.DropDown).GetCurrentOption()
136+
_, action := t.form.GetFormItem(3).(*tview.DropDown).GetCurrentOption()
137+
from := t.form.GetFormItem(4).(*tview.InputField).GetText()
138+
comment := t.form.GetFormItem(5).(*tview.InputField).GetText()
139+
140+
baseCmd := "ufw "
141+
if len(position) > 0 && position[0] < t.table.GetRowCount()-1 {
142+
baseCmd = fmt.Sprintf("ufw insert %d ", position[0])
143+
}
144+
145+
if port == "" {
146+
return
147+
}
148+
149+
cmd := ""
150+
if from == "" && to == "" {
151+
cmd = fmt.Sprintf("%s proto %s to any port %s comment '%s'", strings.ToLower(action), proto, port, comment)
152+
}
153+
154+
if from == "" && to != "" {
155+
cmd = fmt.Sprintf("%s proto %s to %s port %s comment '%s'", strings.ToLower(action), proto, to, port, comment)
156+
}
157+
158+
if to == "" && from != "" {
159+
cmd = fmt.Sprintf("%s from %s proto %s to any port %s comment '%s'", strings.ToLower(action), from, proto, port, comment)
160+
}
161+
162+
if to != "" && from != "" {
163+
cmd = fmt.Sprintf("%s from %s proto %s to %s port %s comment '%s'", strings.ToLower(action), from, proto, to, port, comment)
164+
}
165+
166+
err, _, _ := shellout(baseCmd + cmd)
167+
if err != nil {
168+
log.Print(err)
169+
}
170+
t.Reset()
171+
t.ReloadTable()
172+
}
173+
174+
func (t *Tui) RemoveRule() {
175+
t.table.SetSelectedFunc(func(row int, column int) {
109176
t.table.SetSelectable(false, false)
110177
if row == 0 {
111178
t.app.SetFocus(t.table)
@@ -127,24 +194,70 @@ func (t *Tui) CreateTable(rows []string) {
127194
})
128195
}
129196

130-
func (t *Tui) ReloadTable() {
131-
t.table.Clear()
132-
data, _ := t.LoadTableData()
133-
t.CreateTable(data)
134-
}
197+
func (t *Tui) EditForm() {
198+
t.table.SetSelectedFunc(func(row int, column int) {
199+
if row == 0 {
200+
t.app.SetFocus(t.table)
201+
return
202+
}
203+
t.help.SetText("Use <Tab> and <Enter> keys to navigate through the form").SetBorderPadding(1, 0, 1, 1)
135204

136-
func (t *Tui) CreateModal(text string, confirm func(), cancel func(), finally func()) {
137-
modal := tview.NewModal()
138-
t.pages.AddPage("modal", modal.SetText(text).AddButtons([]string{"Confirm", "Cancel"}).SetDoneFunc(func(i int, label string) {
139-
if label == "Confirm" {
140-
confirm()
141-
t.ReloadTable()
142-
} else {
143-
cancel()
205+
to := t.table.GetCell(row, 1).Text
206+
re := regexp.MustCompile(`((([0-9]{1,3}\.){3}[0-9]{1,3})/)?([a-z]{3})`)
207+
match := re.FindStringSubmatch(to)
208+
209+
toValue := ""
210+
proto := match[0]
211+
if match[4] != "" {
212+
toValue = match[2]
213+
proto = match[4]
144214
}
145-
modal.ClearButtons()
146-
finally()
147-
}), true, true)
215+
216+
protocolOptionIndex := 1
217+
if proto == "tcp" {
218+
protocolOptionIndex = 0
219+
}
220+
221+
actionOptionIndex := 0
222+
switch t.table.GetCell(row, 3).Text {
223+
case "ALLOW":
224+
actionOptionIndex = 0
225+
case "DENY":
226+
actionOptionIndex = 1
227+
case "REJECT":
228+
actionOptionIndex = 2
229+
case "LIMIT":
230+
actionOptionIndex = 3
231+
}
232+
233+
from := t.table.GetCell(row, 4).Text
234+
fromValue := from
235+
if t.table.GetCell(row, 4).Text == "Anywhere" {
236+
fromValue = ""
237+
}
238+
comment := strings.ReplaceAll(t.table.GetCell(row, 5).Text, "# ", "")
239+
240+
t.form.AddInputField("To", toValue, 20, nil, nil).SetFieldTextColor(tcell.ColorWhite).
241+
AddInputField("Port *", t.table.GetCell(row, 2).Text, 20, nil, nil).SetFieldTextColor(tcell.ColorWhite).
242+
AddDropDown("Protocol *", []string{"tcp", "udp"}, protocolOptionIndex, nil).
243+
AddDropDown("Action *", []string{"ALLOW", "DENY", "REJECT", "LIMIT"}, actionOptionIndex, nil).
244+
AddInputField("From", fromValue, 20, nil, nil).
245+
AddInputField("Comment", comment, 40, nil, nil).
246+
AddButton("Save", func() { t.CreateRule(row); t.table.SetSelectable(false, false) }).
247+
AddButton("Cancel", func() { t.Cancel(); t.table.SetSelectable(false, false) }).
248+
SetButtonTextColor(tcell.ColorWhite).
249+
SetButtonBackgroundColor(tcell.ColorDarkCyan).
250+
SetFieldBackgroundColor(tcell.ColorDarkCyan).
251+
SetLabelColor(tcell.ColorWhite)
252+
253+
shellout(fmt.Sprintf("ufw --force delete %d", row))
254+
255+
t.secondHelp.SetText("* Mandatory field\n\nTo and From fields match any and Anywhere if left empty").
256+
SetTextColor(tcell.ColorDarkCyan).
257+
SetBorderPadding(0, 0, 1, 1)
258+
259+
t.app.SetFocus(t.form)
260+
})
148261
}
149262

150263
func (t *Tui) CreateMenu() {
@@ -154,7 +267,12 @@ func (t *Tui) CreateMenu() {
154267
t.CreateForm()
155268
t.app.SetFocus(t.form)
156269
}).
270+
AddItem("Edit a rule", "", 'e', func() {
271+
t.EditForm()
272+
t.app.SetFocus(t.table)
273+
}).
157274
AddItem("Remove a rule", "", 'd', func() {
275+
t.RemoveRule()
158276
t.app.SetFocus(t.table)
159277
t.help.SetText("Press <Esc> to go back to the menu selection").SetBorderPadding(1, 0, 1, 0)
160278
}).
@@ -203,7 +321,7 @@ func (t *Tui) CreateForm() {
203321
AddDropDown("Action *", []string{"ALLOW", "DENY", "REJECT", "LIMIT"}, 0, nil).
204322
AddInputField("From", "", 20, nil, nil).
205323
AddInputField("Comment", "", 40, nil, nil).
206-
AddButton("Save", t.CreateRule).
324+
AddButton("Save", func() { t.CreateRule() }).
207325
AddButton("Cancel", t.Cancel).
208326
SetButtonTextColor(tcell.ColorWhite).
209327
SetButtonBackgroundColor(tcell.ColorDarkCyan).
@@ -221,40 +339,6 @@ func (t *Tui) Reset() {
221339
t.app.SetFocus(t.menu)
222340
}
223341

224-
func (t *Tui) CreateRule() {
225-
to := t.form.GetFormItem(0).(*tview.InputField).GetText()
226-
toPort := t.form.GetFormItem(1).(*tview.InputField).GetText()
227-
_, proto := t.form.GetFormItem(2).(*tview.DropDown).GetCurrentOption()
228-
_, action := t.form.GetFormItem(3).(*tview.DropDown).GetCurrentOption()
229-
from := t.form.GetFormItem(4).(*tview.InputField).GetText()
230-
comment := t.form.GetFormItem(5).(*tview.InputField).GetText()
231-
232-
if toPort == "" {
233-
return
234-
}
235-
236-
cmd := ""
237-
238-
if from == "" && to == "" {
239-
cmd = fmt.Sprintf("ufw %s proto %s to any port %s comment '%s'", strings.ToLower(action), proto, toPort, comment)
240-
}
241-
242-
if from == "" && to != "" {
243-
cmd = fmt.Sprintf("ufw %s proto %s to %s port %s comment '%s'", strings.ToLower(action), proto, to, toPort, comment)
244-
}
245-
246-
if to == "" && from != "" {
247-
cmd = fmt.Sprintf("ufw %s from %s proto %s to any port %s comment '%s'", strings.ToLower(action), from, proto, toPort, comment)
248-
}
249-
250-
err, _, _ := shellout(cmd)
251-
if err != nil {
252-
log.Print(err)
253-
}
254-
t.Reset()
255-
t.ReloadTable()
256-
}
257-
258342
func (t *Tui) Cancel() {
259343
t.Reset()
260344
}

0 commit comments

Comments
 (0)