main
6d86999 ยท 6 months ago 26 commits
  1package ufw
  2
  3import (
  4	"bytes"
  5	"fmt"
  6	"log"
  7	"net"
  8	"os/exec"
  9	"regexp"
 10	"strings"
 11)
 12
 13type Stats struct {
 14	Active     bool
 15	Logging    string
 16	TotalRules int
 17}
 18
 19type Policy struct {
 20	DefaultIncoming string
 21	DefaultOutgoing string
 22	DefaultRouted   string
 23}
 24
 25type Rule struct {
 26	Num    int
 27	Action string
 28	IPv6   bool 
 29
 30	ToDest     string
 31	ToPort     string
 32	ToProtocol string
 33
 34	FromSource string
 35	FromPort   string
 36
 37	Comment string
 38	Raw     string
 39}
 40
 41type ufwData struct {
 42	Stats     Stats
 43	Policy    Policy
 44	Rules     []Rule 
 45	IPv4Rules []Rule
 46	IPv6Rules []Rule
 47	Error     error
 48}
 49
 50// UFW action constants
 51const (
 52	ActionAllow  = "ALLOW"
 53	ActionDeny   = "DENY"
 54	ActionReject = "REJECT"
 55	ActionLimit  = "LIMIT"
 56)
 57
 58var Actions = []string{ActionAllow, ActionDeny, ActionReject, ActionLimit}
 59
 60var (
 61	reStatus  = regexp.MustCompile(`(?i)Status:\s*(\w+)`)
 62	reDefault = regexp.MustCompile(`(?im)^Default:\s*(.+)$`)
 63	reLogging = regexp.MustCompile(`(?im)^Logging:\s*(.+)$`)
 64)
 65
 66func RunCmd(name string, args ...string) (stdout, stderr string, err error) {
 67	cmd := exec.Command(name, args...)
 68	var out, errOut bytes.Buffer
 69	cmd.Stdout = &out
 70	cmd.Stderr = &errOut
 71	err = cmd.Run()
 72	return out.String(), errOut.String(), err
 73}
 74
 75func RunSudo(args ...string) (stdout, stderr string, err error) {
 76	return RunCmd("sudo", append([]string{"ufw"}, args...)...)
 77}
 78
 79func extractPolicy(s string) string {
 80	parts := strings.Fields(s)
 81	if len(parts) > 0 {
 82		return parts[0]
 83	}
 84	return "unknown"
 85}
 86
 87func GetUFWData() ufwData {
 88	var data ufwData
 89
 90	numOut, _, errNum := RunSudo("status", "numbered")
 91	verbOut, _, errVerb := RunSudo("status", "verbose")
 92
 93	if errNum != nil && errVerb != nil {
 94		data.Error = fmt.Errorf("ufw status: %w", errNum)
 95		return data
 96	}
 97
 98	if m := reStatus.FindStringSubmatch(verbOut); len(m) > 1 {
 99		data.Stats.Active = strings.EqualFold(m[1], "active")
100	}
101
102	if m := reDefault.FindStringSubmatch(verbOut); len(m) > 1 {
103		parts := strings.SplitSeq(strings.TrimSpace(m[1]), ",")
104		for p := range parts {
105			p = strings.TrimSpace(p)
106			lower := strings.ToLower(p)
107			if strings.Contains(lower, "incoming") {
108				data.Policy.DefaultIncoming = strings.ToUpper(extractPolicy(p))
109			}
110			if strings.Contains(lower, "outgoing") {
111				data.Policy.DefaultOutgoing = strings.ToUpper(extractPolicy(p))
112			}
113			if strings.Contains(lower, "routed") {
114				data.Policy.DefaultRouted = strings.ToUpper(extractPolicy(p))
115			}
116		}
117	}
118	if m := reLogging.FindStringSubmatch(verbOut); len(m) > 1 {
119		data.Stats.Logging = m[1]
120	}
121
122	data.Rules = ParseRules(numOut)
123	data.Stats.TotalRules = len(data.Rules)
124
125	for _, rule := range data.Rules {
126		if rule.IPv6 {
127			data.IPv6Rules = append(data.IPv6Rules, rule)
128		} else {
129			data.IPv4Rules = append(data.IPv4Rules, rule)
130		}
131	}
132
133	return data
134}
135
136func Enable() (stdout, stderr string, err error) {
137	return RunSudo("enable")
138}
139
140func Disable() (stdout, stderr string, err error) {
141	return RunSudo("disable")
142}
143
144func DeleteRule(num int) (stdout, stderr string, err error) {
145	return RunSudo("--force", "delete", fmt.Sprintf("%d", num))
146}
147
148func DefaultIncoming(allow bool) (stdout, stderr string, err error) {
149	pol := "deny"
150	if allow {
151		pol = "allow"
152	}
153	return RunSudo("default", pol, "incoming")
154}
155
156func DefaultOutgoing(allow bool) (stdout, stderr string, err error) {
157	pol := "deny"
158	if allow {
159		pol = "allow"
160	}
161	return RunSudo("default", pol, "outgoing")
162}
163
164func DefaultRouted(allow bool) (stdout, stderr string, err error) {
165	pol := "deny"
166	if allow {
167		pol = "allow"
168	}
169	return RunSudo("default", pol, "routed")
170}
171
172func SetLogging(level string) (stdout, stderr string, err error) {
173	return RunSudo("logging", level)
174}
175
176func AddRule(rule string) (stdout, stderr string, err error) {
177	parts := strings.Fields(rule)
178	if len(parts) == 0 {
179		return "", "", fmt.Errorf("empty rule")
180	}
181	return RunSudo(append([]string{}, parts...)...)
182}
183
184func InsertRule(position int, args ...string) (stdout, stderr string, err error) {
185	cmdArgs := append([]string{"insert", fmt.Sprintf("%d", position)}, args...)
186	log.Printf("InsertRule cmdArgs: %v", cmdArgs)
187	return RunSudo(cmdArgs...)
188}
189
190// normalizeAddress cleans up UFW display addresses for use in commands
191// "Anywhere (v6)" -> "" (let UFW handle it)
192// "Anywhere" -> "" (let UFW handle it)
193// "192.168.1.0/24" -> "192.168.1.0/24"
194func normalizeAddress(addr string) string {
195	addr = strings.TrimSpace(addr)
196	addr = strings.TrimSuffix(addr, " (v6)")
197	if addr == "Anywhere" || addr == "" {
198		return ""
199	}
200	return addr
201}
202
203// normalizePort cleans up UFW display port values for use in commands
204// "80 (v6)" -> "80"
205// "22/tcp (v6)" -> "22/tcp"
206func normalizePort(port string) string {
207	port = strings.TrimSpace(port)
208	port = strings.TrimSuffix(port, " (v6)")
209	return port
210}
211
212func InsertRuleFromExisting(position int, action string, rule Rule) (stdout, stderr string, err error) {
213	args := buildRuleArgs(action, rule)
214
215	log.Printf("InsertRuleFromExisting: position=%d, action=%s, ipv6=%v, args=%v", position, action, rule.IPv6, args)
216
217	if position <= 0 {
218		return AppendRule(args...)
219	}
220	return InsertRule(position, args...)
221}
222
223func AppendRule(args ...string) (stdout, stderr string, err error) {
224	log.Printf("AppendRule args: %v", args)
225	return RunSudo(args...)
226}
227
228func buildRuleArgs(action string, rule Rule) []string {
229	var args []string
230	args = append(args, strings.ToLower(action))
231
232	toDest := normalizeAddress(rule.ToDest)
233	toPort := normalizePort(rule.ToPort)
234	toProto := rule.ToProtocol
235
236	fromSource := normalizeAddress(rule.FromSource)
237	fromPort := normalizePort(rule.FromPort)
238
239	// For IPv6 rules, we need to use specific addresses or "any" will default to IPv4
240	// UFW creates IPv6 rules when using IPv6 addresses or when the original rule was IPv6
241	// We use "::/0" for IPv6 "anywhere" equivalent
242	anyAddr := "any"
243	if rule.IPv6 {
244		anyAddr = "::/0" // IPv6 equivalent of "anywhere"
245	}
246
247	if fromSource != "" {
248		args = append(args, "from", fromSource)
249	} else {
250		args = append(args, "from", anyAddr)
251	}
252
253	if toDest != "" {
254		args = append(args, "to", toDest)
255	} else {
256		args = append(args, "to", anyAddr)
257	}
258
259	if toPort != "" && toPort != "any" {
260		args = append(args, "port", toPort)
261	}
262
263	if toProto != "" && toProto != "any" {
264		args = append(args, "proto", strings.ToLower(toProto))
265	}
266
267	if fromPort != "" && fromPort != "any" {
268		args = append(args, "sport", fromPort)
269	}
270
271	return args
272}
273
274func GetCurrentRules() []Rule {
275	numOut, _, err := RunSudo("status", "numbered")
276	if err != nil {
277		log.Printf("Error fetching rules: %v", err)
278		return nil
279	}
280	return ParseRules(numOut)
281}
282
283func FindInsertPosition(rules []Rule, prevNextNum int, isIPv6 bool) int {
284	if prevNextNum <= 0 {
285		return 0
286	}
287
288	for _, r := range rules {
289		if r.IPv6 == isIPv6 && r.Num >= 1 {
290			return r.Num
291		}
292	}
293	return 0
294}
295
296func MoveRule(rule Rule, direction int, newAction string) error {
297	originalNum := rule.Num
298	isIPv6 := rule.IPv6
299
300	log.Printf("MoveRule: rule #%d, direction=%d, newAction=%s, isIPv6=%v", originalNum, direction, newAction, isIPv6)
301
302	// Step 1: Get current rules to understand the landscape
303	currentRules := GetCurrentRules()
304	if currentRules == nil {
305		return fmt.Errorf("failed to fetch current rules")
306	}
307
308	// Find our rule's neighbors of the same IP version
309	var sameTypeRules []Rule
310	for _, r := range currentRules {
311		if r.IPv6 == isIPv6 {
312			sameTypeRules = append(sameTypeRules, r)
313		}
314	}
315
316	// Find index of our rule in the same-type list
317	ourIndex := -1
318	for i, r := range sameTypeRules {
319		if r.Num == originalNum {
320			ourIndex = i
321			break
322		}
323	}
324
325	if ourIndex == -1 {
326		return fmt.Errorf("rule #%d not found in current rules", originalNum)
327	}
328
329	// Calculate target index in same-type list
330	targetIndex := max(ourIndex + direction, 0)
331	if targetIndex >= len(sameTypeRules) {
332		targetIndex = len(sameTypeRules) - 1
333	}
334
335	// Step 2: Delete the rule
336	_, stderr, err := DeleteRule(originalNum)
337	if err != nil {
338		return fmt.Errorf("failed to delete rule: %s", stderr)
339	}
340
341	// Step 3: Re-fetch rules to get accurate positions
342	newRules := GetCurrentRules()
343	if newRules == nil {
344		return fmt.Errorf("failed to fetch rules after delete")
345	}
346
347	// Find same-type rules again after deletion
348	var newSameTypeRules []Rule
349	for _, r := range newRules {
350		if r.IPv6 == isIPv6 {
351			newSameTypeRules = append(newSameTypeRules, r)
352		}
353	}
354
355	// Step 4: Calculate actual insert position
356	var insertPos int
357	if len(newSameTypeRules) == 0 {
358		insertPos = 0
359	} else if targetIndex >= len(newSameTypeRules) {
360		insertPos = 0
361	} else {
362		insertPos = newSameTypeRules[targetIndex].Num
363	}
364
365	log.Printf("MoveRule: targetIndex=%d, insertPos=%d, newSameTypeRules=%d", targetIndex, insertPos, len(newSameTypeRules))
366
367	// Step 5: Validate insert position
368	// UFW insert position must be between 1 and (total_rules + 1)
369	totalRulesAfterDelete := len(newRules)
370	if insertPos > totalRulesAfterDelete+1 {
371		log.Printf("MoveRule: insertPos %d exceeds valid range (1-%d), appending instead", insertPos, totalRulesAfterDelete+1)
372		insertPos = 0 // Signal to append
373	}
374
375	// Step 6: Insert the rule at the calculated position
376	_, stderr, err = InsertRuleFromExisting(insertPos, newAction, rule)
377	if err != nil {
378		return fmt.Errorf("failed to insert rule (pos=%d, total=%d): %s", insertPos, totalRulesAfterDelete, stderr)
379	}
380
381	return nil
382}
383
384func GetInterfaces() []string {
385	interfaces, err := net.Interfaces()
386	if err != nil {
387		log.Printf("Error getting interfaces: %v", err)
388		return []string{}
389	}
390
391	var names []string
392	for _, iface := range interfaces {
393		if iface.Flags&net.FlagLoopback != 0 {
394			continue
395		}
396		names = append(names, iface.Name)
397	}
398	return names
399}
400
401const (
402	DirectionIn  = "in"
403	DirectionOut = "out"
404)
405
406const (
407	ProtoTCP = "tcp"
408	ProtoUDP = "udp"
409	ProtoAny = "any"
410)
411
412var Directions = []string{DirectionIn, DirectionOut}
413var Protocols = []string{ProtoAny, ProtoTCP, ProtoUDP}
414var CommonPorts = []struct {
415	Port string
416	Name string
417}{
418	{"22", "SSH"},
419	{"80", "HTTP"},
420	{"443", "HTTPS"},
421	{"53", "DNS"},
422	{"5432", "PostgreSQL"},
423	{"6379", "Redis"},
424	{"3000", "Dev Server"},
425}
426
427type AddRuleParams struct {
428	Action    string // allow, deny, reject, limit
429	Direction string // in, out, or empty for both
430	Protocol  string // tcp, udp, any
431	Port      string // port number or range
432	FromAddr  string // source address or empty for any
433	ToAddr    string // destination address or empty for any
434	Interface string // network interface or empty for all
435	Comment   string // optional comment
436}
437
438func BuildAddRuleCommand(p AddRuleParams) []string {
439	var args []string
440
441	args = append(args, strings.ToLower(p.Action))
442
443	if p.Direction != "" {
444		args = append(args, p.Direction)
445	}
446
447	if p.Interface != "" {
448		args = append(args, "on", p.Interface)
449	}
450
451	if p.Port != "" {
452		if p.FromAddr != "" && p.FromAddr != "any" {
453			args = append(args, "from", p.FromAddr)
454		} else {
455			args = append(args, "from", "any")
456		}
457
458		if p.ToAddr != "" && p.ToAddr != "any" {
459			args = append(args, "to", p.ToAddr)
460		} else {
461			args = append(args, "to", "any")
462		}
463
464		args = append(args, "port", p.Port)
465
466		if p.Protocol != "" && p.Protocol != ProtoAny {
467			args = append(args, "proto", p.Protocol)
468		}
469	} else {
470		if p.FromAddr != "" && p.FromAddr != "any" {
471			args = append(args, "from", p.FromAddr)
472		}
473		if p.ToAddr != "" && p.ToAddr != "any" {
474			args = append(args, "to", p.ToAddr)
475		}
476	}
477
478	if p.Comment != "" {
479		args = append(args, "comment", p.Comment)
480	}
481
482	return args
483}
484
485func AddNewRule(p AddRuleParams) (stdout, stderr string, err error) {
486	args := BuildAddRuleCommand(p)
487	log.Printf("AddNewRule: %v", args)
488	return RunSudo(args...)
489}