Skip to main content

jpx_core/extensions/
network.rs

1//! Network and IP address functions.
2
3use 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
15// =============================================================================
16// ip_to_int(s) -> number
17// =============================================================================
18
19defn!(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
36// =============================================================================
37// int_to_ip(n) -> string
38// =============================================================================
39
40defn!(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
56// =============================================================================
57// cidr_contains(cidr, ip) -> bool
58// =============================================================================
59
60defn!(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
82// =============================================================================
83// cidr_network(cidr) -> string
84// =============================================================================
85
86defn!(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
100// =============================================================================
101// cidr_broadcast(cidr) -> string
102// =============================================================================
103
104defn!(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
118// =============================================================================
119// cidr_prefix(cidr) -> number
120// =============================================================================
121
122defn!(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
136// =============================================================================
137// is_private_ip(ip) -> bool
138// =============================================================================
139
140defn!(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
154/// Register network functions filtered by the enabled set.
155pub 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        // 192.168.1.1 = 192*256^3 + 168*256^2 + 1*256 + 1 = 3232235777
209        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        // 192.168.x.x is private
279        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        // 10.x.x.x is private
289        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        // 8.8.8.8 is public (Google DNS)
299        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}