From c642b5eeae186f81968cdad6f93d3404c590314b Mon Sep 17 00:00:00 2001
From: James Shubin <james@shubin.ca>
Date: Wed, 6 Nov 2024 21:10:04 -0500
Subject: [PATCH] lang: core: net: Add new function to get cidr prefix

---
 lang/core/net/cidr_to_ip_func.go | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/lang/core/net/cidr_to_ip_func.go b/lang/core/net/cidr_to_ip_func.go
index cb42e2a0..94cea903 100644
--- a/lang/core/net/cidr_to_ip_func.go
+++ b/lang/core/net/cidr_to_ip_func.go
@@ -32,6 +32,7 @@ package corenet
 import (
 	"context"
 	"net"
+	"strconv"
 	"strings"
 
 	"github.com/purpleidea/mgmt/lang/funcs/simple"
@@ -43,6 +44,10 @@ func init() {
 		T: types.NewType("func(a str) str"),
 		F: CidrToIP,
 	})
+	simple.ModuleRegister(ModuleName, "cidr_to_prefix", &simple.Scaffold{
+		T: types.NewType("func(a str) str"),
+		F: CidrToPrefix,
+	})
 	simple.ModuleRegister(ModuleName, "cidr_to_mask", &simple.Scaffold{
 		T: types.NewType("func(a str) str"),
 		F: CidrToMask,
@@ -61,6 +66,22 @@ func CidrToIP(ctx context.Context, input []types.Value) (types.Value, error) {
 	}, nil
 }
 
+// CidrToPrefix returns the prefix from a CIDR address. For example, if you give
+// us 192.0.2.0/24 then we will return "24" as a string.
+func CidrToPrefix(ctx context.Context, input []types.Value) (types.Value, error) {
+	cidr := input[0].Str()
+	_, ipnet, err := net.ParseCIDR(strings.TrimSpace(cidr))
+	if err != nil {
+		return nil, err
+	}
+
+	ones, _ := ipnet.Mask.Size()
+
+	return &types.StrValue{
+		V: strconv.Itoa(ones),
+	}, nil
+}
+
 // CidrToMask returns the subnet mask from a CIDR address.
 func CidrToMask(ctx context.Context, input []types.Value) (types.Value, error) {
 	cidr := input[0].Str()