From f47bf69ee1b58c466e870d2d71284a518a1e9378 Mon Sep 17 00:00:00 2001
From: Craig Johnston <cj@imti.co>
Date: Sun, 23 Jul 2023 19:54:05 -0700
Subject: [PATCH] added RemoveCIDRs

---
 txeh.go                 | 40 ++++++++++++++++++++++++++++++++++++++--
 util/cmd/remove_cidr.go | 35 ++++++-----------------------------
 2 files changed, 44 insertions(+), 31 deletions(-)

diff --git a/txeh.go b/txeh.go
index 0b62864..0744bc6 100644
--- a/txeh.go
+++ b/txeh.go
@@ -111,7 +111,7 @@ func (h *Hosts) Save() error {
 // SaveAs saves rendered hosts file to the filename specified
 func (h *Hosts) SaveAs(fileName string) error {
 	if h.RawText != nil {
-		return errors.New("can not call Save or SaveAs with RawText. Use RenderHostsFile to return a string")
+		return errors.New("cannot call Save or SaveAs with RawText. Use RenderHostsFile to return a string")
 	}
 	hfData := []byte(h.RenderHostsFile())
 
@@ -128,6 +128,9 @@ func (h *Hosts) SaveAs(fileName string) error {
 
 // Reload hosts file
 func (h *Hosts) Reload() error {
+	if h.RawText != nil {
+		return errors.New("cannot call Reload with RawText")
+	}
 	h.Lock()
 	defer h.Unlock()
 
@@ -157,7 +160,7 @@ func (h *Hosts) RemoveAddress(address string) {
 	}
 }
 
-// RemoveFirstAddress removed the first entry (line) found with the provided address.
+// RemoveFirstAddress removes the first entry (line) found with the provided address.
 func (h *Hosts) RemoveFirstAddress(address string) bool {
 	h.Lock()
 	defer h.Unlock()
@@ -172,6 +175,39 @@ func (h *Hosts) RemoveFirstAddress(address string) bool {
 	return false
 }
 
+// RemoveCIDRs Remove CIDR Range (Classless inter-domain routing)
+// examples:
+//
+//	127.1.0.0/16  = 127.1.0.0  -> 127.1.255.255
+//	127.1.27.0/24 = 127.1.27.0 -> 127.1.27.255
+func (h *Hosts) RemoveCIDRs(cidrs []string) error {
+	addresses := make([]string, 0)
+
+	// loop through all the CIDR ranges (we probably have less ranges than IPs)
+	for _, cidr := range cidrs {
+
+		_, ipnet, err := net.ParseCIDR(cidr)
+		if err != nil {
+			return err
+		}
+
+		hfLines := h.GetHostFileLines()
+
+		for _, hfl := range *hfLines {
+			ip := net.ParseIP(hfl.Address)
+			if ip != nil {
+				if ipnet.Contains(ip) {
+					addresses = append(addresses, hfl.Address)
+				}
+			}
+		}
+	}
+
+	h.RemoveAddresses(addresses)
+
+	return nil
+}
+
 // RemoveHosts removes all hostname entries of the provided host slice
 func (h *Hosts) RemoveHosts(hosts []string) {
 	for _, host := range hosts {
diff --git a/util/cmd/remove_cidr.go b/util/cmd/remove_cidr.go
index 78ad7d0..7a94160 100644
--- a/util/cmd/remove_cidr.go
+++ b/util/cmd/remove_cidr.go
@@ -3,11 +3,9 @@ package cmd
 import (
 	"errors"
 	"fmt"
-	"net"
+	"github.com/spf13/cobra"
 	"os"
 	"strings"
-
-	"github.com/spf13/cobra"
 )
 
 func init() {
@@ -40,39 +38,18 @@ var removeCidrCmd = &cobra.Command{
 
 func RemoveIPRanges(cidrs []string) {
 
-	addresses := make([]string, 0)
-
-	// loop through all the CIDR ranges
-	for _, cidr := range cidrs {
-
-		_, ipnet, err := net.ParseCIDR(cidr)
-		if err != nil {
-			fmt.Printf("Error: there was a problem with the IP range %s. Reason: %s\n", cidr, err.Error())
-			os.Exit(1)
-		}
-
-		hfLines := etcHosts.GetHostFileLines()
-
-		for _, hfl := range *hfLines {
-
-			ip := net.ParseIP(hfl.Address)
-			if ip != nil {
-				if ipnet.Contains(ip) {
-					addresses = append(addresses, hfl.Address)
-				}
-			}
-		}
-
+	err := etcHosts.RemoveCIDRs(cidrs)
+	if err != nil {
+		fmt.Printf("Error: there was a problem parsing a CIDR. Reason: %s\n", err.Error())
+		os.Exit(1)
 	}
 
-	etcHosts.RemoveAddresses(addresses)
-
 	if DryRun {
 		fmt.Print(etcHosts.RenderHostsFile())
 		return
 	}
 
-	err := etcHosts.Save()
+	err = etcHosts.Save()
 	if err != nil {
 		fmt.Printf("Error: could not save %s. Reason: %s\n", etcHosts.WriteFilePath, err.Error())
 		os.Exit(1)