1use std::collections::HashSet;
4use std::net::Ipv4Addr;
5use std::str::FromStr;
6
7use ipnetwork::{IpNetwork, Ipv4Network};
8use serde_json::{Number, Value};
9
10use crate::functions::{Function, number_value};
11use crate::interpreter::SearchResult;
12use crate::registry::register_if_enabled;
13use crate::{Context, Runtime, arg, defn};
14
15defn!(IpToIntFn, vec![arg!(string)], None);
20
21impl Function for IpToIntFn {
22 fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
23 self.signature.validate(args, ctx)?;
24 let s = args[0].as_str().unwrap();
25
26 match Ipv4Addr::from_str(s) {
27 Ok(ip) => {
28 let int_val: u32 = ip.into();
29 Ok(number_value(int_val as f64))
30 }
31 Err(_) => Ok(Value::Null),
32 }
33 }
34}
35
36defn!(IntToIpFn, vec![arg!(number)], None);
41
42impl Function for IntToIpFn {
43 fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
44 self.signature.validate(args, ctx)?;
45 let n = args[0].as_f64().unwrap();
46
47 if n < 0.0 || n > u32::MAX as f64 {
48 return Ok(Value::Null);
49 }
50
51 let ip = Ipv4Addr::from(n as u32);
52 Ok(Value::String(ip.to_string()))
53 }
54}
55
56defn!(CidrContainsFn, vec![arg!(string), arg!(string)], None);
61
62impl Function for CidrContainsFn {
63 fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
64 self.signature.validate(args, ctx)?;
65 let cidr_str = args[0].as_str().unwrap();
66 let ip_str = args[1].as_str().unwrap();
67
68 let network = match IpNetwork::from_str(cidr_str) {
69 Ok(n) => n,
70 Err(_) => return Ok(Value::Null),
71 };
72
73 let ip: std::net::IpAddr = match ip_str.parse() {
74 Ok(ip) => ip,
75 Err(_) => return Ok(Value::Null),
76 };
77
78 Ok(Value::Bool(network.contains(ip)))
79 }
80}
81
82defn!(CidrNetworkFn, vec![arg!(string)], None);
87
88impl Function for CidrNetworkFn {
89 fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
90 self.signature.validate(args, ctx)?;
91 let cidr_str = args[0].as_str().unwrap();
92
93 match Ipv4Network::from_str(cidr_str) {
94 Ok(network) => Ok(Value::String(network.network().to_string())),
95 Err(_) => Ok(Value::Null),
96 }
97 }
98}
99
100defn!(CidrBroadcastFn, vec![arg!(string)], None);
105
106impl Function for CidrBroadcastFn {
107 fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
108 self.signature.validate(args, ctx)?;
109 let cidr_str = args[0].as_str().unwrap();
110
111 match Ipv4Network::from_str(cidr_str) {
112 Ok(network) => Ok(Value::String(network.broadcast().to_string())),
113 Err(_) => Ok(Value::Null),
114 }
115 }
116}
117
118defn!(CidrPrefixFn, vec![arg!(string)], None);
123
124impl Function for CidrPrefixFn {
125 fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
126 self.signature.validate(args, ctx)?;
127 let cidr_str = args[0].as_str().unwrap();
128
129 match IpNetwork::from_str(cidr_str) {
130 Ok(network) => Ok(Value::Number(Number::from(network.prefix()))),
131 Err(_) => Ok(Value::Null),
132 }
133 }
134}
135
136defn!(IsPrivateIpFn, vec![arg!(string)], None);
141
142impl Function for IsPrivateIpFn {
143 fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
144 self.signature.validate(args, ctx)?;
145 let ip_str = args[0].as_str().unwrap();
146
147 match Ipv4Addr::from_str(ip_str) {
148 Ok(ip) => Ok(Value::Bool(ip.is_private())),
149 Err(_) => Ok(Value::Null),
150 }
151 }
152}
153
154pub fn register_filtered(runtime: &mut Runtime, enabled: &HashSet<&str>) {
156 register_if_enabled(runtime, "ip_to_int", enabled, Box::new(IpToIntFn::new()));
157 register_if_enabled(runtime, "int_to_ip", enabled, Box::new(IntToIpFn::new()));
158 register_if_enabled(
159 runtime,
160 "cidr_contains",
161 enabled,
162 Box::new(CidrContainsFn::new()),
163 );
164 register_if_enabled(
165 runtime,
166 "cidr_network",
167 enabled,
168 Box::new(CidrNetworkFn::new()),
169 );
170 register_if_enabled(
171 runtime,
172 "cidr_broadcast",
173 enabled,
174 Box::new(CidrBroadcastFn::new()),
175 );
176 register_if_enabled(
177 runtime,
178 "cidr_prefix",
179 enabled,
180 Box::new(CidrPrefixFn::new()),
181 );
182 register_if_enabled(
183 runtime,
184 "is_private_ip",
185 enabled,
186 Box::new(IsPrivateIpFn::new()),
187 );
188}
189
190#[cfg(test)]
191mod tests {
192 use crate::Runtime;
193 use serde_json::json;
194
195 fn setup_runtime() -> Runtime {
196 Runtime::builder()
197 .with_standard()
198 .with_all_extensions()
199 .build()
200 }
201
202 #[test]
203 fn test_ip_to_int() {
204 let runtime = setup_runtime();
205 let data = json!("192.168.1.1");
206 let expr = runtime.compile("ip_to_int(@)").unwrap();
207 let result = expr.search(&data).unwrap();
208 assert_eq!(result.as_f64().unwrap(), 3232235777.0);
210 }
211
212 #[test]
213 fn test_int_to_ip() {
214 let runtime = setup_runtime();
215 let data = json!(3232235777_u64);
216 let expr = runtime.compile("int_to_ip(@)").unwrap();
217 let result = expr.search(&data).unwrap();
218 assert_eq!(result.as_str().unwrap(), "192.168.1.1");
219 }
220
221 #[test]
222 fn test_ip_roundtrip() {
223 let runtime = setup_runtime();
224 let data = json!("10.0.0.1");
225 let expr = runtime.compile("int_to_ip(ip_to_int(@))").unwrap();
226 let result = expr.search(&data).unwrap();
227 assert_eq!(result.as_str().unwrap(), "10.0.0.1");
228 }
229
230 #[test]
231 fn test_cidr_contains_true() {
232 let runtime = setup_runtime();
233 let data = json!({"cidr": "192.168.1.0/24", "ip": "192.168.1.100"});
234 let expr = runtime.compile("cidr_contains(cidr, ip)").unwrap();
235 let result = expr.search(&data).unwrap();
236 assert_eq!(result, json!(true));
237 }
238
239 #[test]
240 fn test_cidr_contains_false() {
241 let runtime = setup_runtime();
242 let data = json!({"cidr": "192.168.1.0/24", "ip": "192.168.2.1"});
243 let expr = runtime.compile("cidr_contains(cidr, ip)").unwrap();
244 let result = expr.search(&data).unwrap();
245 assert_eq!(result, json!(false));
246 }
247
248 #[test]
249 fn test_cidr_network() {
250 let runtime = setup_runtime();
251 let data = json!("192.168.1.100/24");
252 let expr = runtime.compile("cidr_network(@)").unwrap();
253 let result = expr.search(&data).unwrap();
254 assert_eq!(result.as_str().unwrap(), "192.168.1.0");
255 }
256
257 #[test]
258 fn test_cidr_broadcast() {
259 let runtime = setup_runtime();
260 let data = json!("192.168.1.0/24");
261 let expr = runtime.compile("cidr_broadcast(@)").unwrap();
262 let result = expr.search(&data).unwrap();
263 assert_eq!(result.as_str().unwrap(), "192.168.1.255");
264 }
265
266 #[test]
267 fn test_cidr_prefix() {
268 let runtime = setup_runtime();
269 let data = json!("10.0.0.0/8");
270 let expr = runtime.compile("cidr_prefix(@)").unwrap();
271 let result = expr.search(&data).unwrap();
272 assert_eq!(result.as_f64().unwrap(), 8.0);
273 }
274
275 #[test]
276 fn test_is_private_ip_true() {
277 let runtime = setup_runtime();
278 let data = json!("192.168.1.1");
280 let expr = runtime.compile("is_private_ip(@)").unwrap();
281 let result = expr.search(&data).unwrap();
282 assert_eq!(result, json!(true));
283 }
284
285 #[test]
286 fn test_is_private_ip_10() {
287 let runtime = setup_runtime();
288 let data = json!("10.0.0.1");
290 let expr = runtime.compile("is_private_ip(@)").unwrap();
291 let result = expr.search(&data).unwrap();
292 assert_eq!(result, json!(true));
293 }
294
295 #[test]
296 fn test_is_private_ip_false() {
297 let runtime = setup_runtime();
298 let data = json!("8.8.8.8");
300 let expr = runtime.compile("is_private_ip(@)").unwrap();
301 let result = expr.search(&data).unwrap();
302 assert_eq!(result, json!(false));
303 }
304}