From b5f81fcc559de48b65f771c2b9dd243e2c2fdb39 Mon Sep 17 00:00:00 2001
From: J$ <jon.monette@gmail.com>
Date: Wed, 12 Jun 2024 11:00:45 -0500
Subject: [PATCH] Add a calculate optino the estimates the file size of the
 encoded file

---
 cmd/cli/main.go                   | 53 ++++++++++++++++++++++---------
 internal/cidrencode/cidrencode.go | 29 +++++++++++------
 2 files changed, 57 insertions(+), 25 deletions(-)

diff --git a/cmd/cli/main.go b/cmd/cli/main.go
index 3c037e6..e2003bb 100644
--- a/cmd/cli/main.go
+++ b/cmd/cli/main.go
@@ -12,38 +12,44 @@ import (
 	"strings"
 
 	"github.com/jmoney/cidr-encoder/internal/cidrencode"
+	"github.com/projectdiscovery/mapcidr"
 )
 
 var (
 	name   = flag.String("name", "", "The file base name to use as the ACL file name(e.g test.acl name is test)")
 	encode = flag.Bool("encode", false, "Encode the CIDRs")
 	search = flag.String("search", "", "Search for a IP in the CIDRs")
+	calc   = flag.Bool("calc", false, "Calculate the size of the encoded file")
 )
 
 func main() {
 	flag.Parse()
-	if *encode {
+
+	if *calc {
 		stdin, err := io.ReadAll(os.Stdin)
 		if err != nil {
 			log.Fatalf("failed to read from stdin: %s", err)
 		}
 		str := string(stdin)
 		cidrs := strings.Split(strings.TrimSuffix(str, "\n"), "\n")
-
-		networks := []*net.IPNet{}
-		for _, cidr := range cidrs {
-
-			if strings.HasPrefix(cidr, "#") || strings.HasPrefix(cidr, "//") || cidr == "" {
-				continue
-			}
-
-			_, network, err := net.ParseCIDR(cidr)
-			if err != nil {
-				log.Printf("Skipping invalid CIDR %s: %s", cidr, err)
-				continue
-			}
-			networks = append(networks, network)
+		min, max, estimatedSize := cidrencode.Calculate(convertStrToNetwork(cidrs))
+		result := map[string]interface{}{
+			"min":           mapcidr.Inet_ntoa(min),
+			"max":           mapcidr.Inet_ntoa(max),
+			"estimatedSize": estimatedSize,
+			"humanReadable": cidrencode.BytesToHumanReadable(estimatedSize),
+		}
+		output, _ := json.Marshal(result)
+		fmt.Println(string(output))
+	} else if *encode {
+		stdin, err := io.ReadAll(os.Stdin)
+		if err != nil {
+			log.Fatalf("failed to read from stdin: %s", err)
 		}
+		str := string(stdin)
+		cidrs := strings.Split(strings.TrimSuffix(str, "\n"), "\n")
+
+		networks := convertStrToNetwork(cidrs)
 
 		file, err := os.OpenFile(filepath.Clean(file(*name)), os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
 		if err != nil {
@@ -76,6 +82,23 @@ func main() {
 	}
 }
 
+func convertStrToNetwork(cidrs []string) []*net.IPNet {
+	networks := []*net.IPNet{}
+	for _, cidr := range cidrs {
+		if strings.HasPrefix(cidr, "#") || strings.HasPrefix(cidr, "//") || cidr == "" {
+			continue
+		}
+
+		_, network, err := net.ParseCIDR(cidr)
+		if err != nil {
+			log.Printf("Skipping invalid CIDR %s: %s", cidr, err)
+			continue
+		}
+		networks = append(networks, network)
+	}
+	return networks
+}
+
 func file(id string) string {
 	return fmt.Sprintf("./%s.acl", id)
 }
diff --git a/internal/cidrencode/cidrencode.go b/internal/cidrencode/cidrencode.go
index 9d17824..67adb6a 100644
--- a/internal/cidrencode/cidrencode.go
+++ b/internal/cidrencode/cidrencode.go
@@ -16,6 +16,13 @@ const (
 	MAGIC_BYTE = byte(1)
 )
 
+func Calculate(networks []*net.IPNet) (int64, int64, int64) {
+	ips := convertToRangePairs(networks)
+	min, max := findMinMax(ips)
+	estimatedFileSize := max - min
+	return min, max, estimatedFileSize
+}
+
 func Search(file *os.File, ip *net.IP) bool {
 	b := make([]byte, 8)
 	file.ReadAt(b, 0)
@@ -27,15 +34,8 @@ func Search(file *os.File, ip *net.IP) bool {
 }
 
 func Encode(file *os.File, networks []*net.IPNet) {
-	ips := []int64{}
-	for _, network := range networks {
-		first, last, _ := mapcidr.AddressRange(network)
-		ips = append(ips, mapcidr.Inet_aton(first), mapcidr.Inet_aton(last))
-	}
-
-	min, max := findMinMax(ips)
-	estimatedFileSize := max - min
-	log.Printf("Estimated File Size: %s\n", bytesToHumanReadable(estimatedFileSize))
+	min, _, estimatedFileSize := Calculate(networks)
+	log.Printf("Estimated File Size: %s\n", BytesToHumanReadable(estimatedFileSize))
 	offset := min - int64(reflect.TypeOf(int64(0)).Size())
 	log.Printf("Offset: %d", offset)
 	b := make([]byte, 8)
@@ -54,7 +54,7 @@ func Encode(file *os.File, networks []*net.IPNet) {
 	}
 }
 
-func bytesToHumanReadable(bytes int64) string {
+func BytesToHumanReadable(bytes int64) string {
 	units := []string{"B", "KB", "MB", "GB", "TB", "PB", "EB"}
 	if bytes < 10 {
 		return fmt.Sprintf("%d B", bytes) // Handle small bytes directly
@@ -81,3 +81,12 @@ func findMinMax(a []int64) (int64, int64) {
 	}
 	return min, max
 }
+
+func convertToRangePairs(networks []*net.IPNet) []int64 {
+	ips := []int64{}
+	for _, network := range networks {
+		first, last, _ := mapcidr.AddressRange(network)
+		ips = append(ips, mapcidr.Inet_aton(first), mapcidr.Inet_aton(last))
+	}
+	return ips
+}