1use crate::compiler::prelude::*;
2use crate::stdlib::ip_utils::to_key;
3use ipcrypt_rs::{Ipcrypt, IpcryptPfx};
4use std::net::IpAddr;
5
6fn encrypt_ip(ip: &Value, key: Value, mode: &Value) -> Resolved {
7 let ip_str = ip.try_bytes_utf8_lossy()?;
8 let ip_addr: IpAddr = ip_str
9 .parse()
10 .map_err(|err| format!("unable to parse IP address: {err}"))?;
11
12 let mode_str = mode.try_bytes_utf8_lossy()?;
13
14 let ip_ver_label = match ip_addr {
15 IpAddr::V4(_) => "IPv4",
16 IpAddr::V6(_) => "IPv6",
17 };
18
19 let encrypted_ip = match mode_str.as_ref() {
20 "aes128" => {
21 let key = to_key::<16>(key, "aes128", ip_ver_label)?;
22 Ipcrypt::new(key).encrypt_ipaddr(ip_addr)
23 }
24 "pfx" => {
25 let key = to_key::<32>(key, "pfx", ip_ver_label)?;
26 IpcryptPfx::new(key).encrypt_ipaddr(ip_addr)
27 }
28 other => {
29 return Err(format!("Invalid mode '{other}'. Must be 'aes128' or 'pfx'").into());
30 }
31 };
32
33 Ok(encrypted_ip.to_string().into())
34}
35
36#[derive(Clone, Copy, Debug)]
37pub struct EncryptIp;
38
39impl Function for EncryptIp {
40 fn identifier(&self) -> &'static str {
41 "encrypt_ip"
42 }
43
44 fn usage(&self) -> &'static str {
45 indoc! {"
46 Encrypts an IP address, transforming it into a different valid IP address.
47
48 Supported Modes:
49
50 * AES128 - Scrambles the entire IP address using AES-128 encryption. Can transform between IPv4 and IPv6.
51 * PFX (Prefix-preserving) - Maintains network hierarchy by ensuring that IP addresses within the same network are encrypted to addresses that also share a common network. This preserves prefix relationships while providing confidentiality.
52 "}
53 }
54
55 fn category(&self) -> &'static str {
56 Category::Ip.as_ref()
57 }
58
59 fn internal_failure_reasons(&self) -> &'static [&'static str] {
60 &[
61 "`ip` is not a valid IP address.",
62 "`mode` is not a supported mode (must be `aes128` or `pfx`).",
63 "`key` length does not match the requirements for the specified mode (16 bytes for `aes128`, 32 bytes for `pfx`).",
64 ]
65 }
66
67 fn return_kind(&self) -> u16 {
68 kind::BYTES
69 }
70
71 fn notices(&self) -> &'static [&'static str] {
72 &[indoc! {"
73 The `aes128` mode implements the `ipcrypt-deterministic` algorithm from the IPCrypt
74 specification, while the `pfx` mode implements the `ipcrypt-pfx` algorithm. Both modes
75 provide deterministic encryption where the same input IP address encrypted with the
76 same key will always produce the same encrypted output.
77 "}]
78 }
79
80 fn parameters(&self) -> &'static [Parameter] {
81 const PARAMETERS: &[Parameter] = &[
82 Parameter::required("ip", kind::BYTES, "The IP address to encrypt (v4 or v6)."),
83 Parameter::required(
84 "key",
85 kind::BYTES,
86 "The encryption key in raw bytes (not encoded). For AES128 mode, the key must be exactly 16 bytes. For PFX mode, the key must be exactly 32 bytes.",
87 ),
88 Parameter::required(
89 "mode",
90 kind::BYTES,
91 "The encryption mode to use. Must be either `aes128` or `pfx`.",
92 ),
93 ];
94 PARAMETERS
95 }
96
97 fn examples(&self) -> &'static [Example] {
98 &[
99 example! {
100 title: "Encrypt IPv4 address with AES128",
101 source: r#"encrypt_ip!("192.168.1.1", "sixteen byte key", "aes128")"#,
102 result: Ok("72b9:a747:f2e9:72af:76ca:5866:6dcf:c3b0"),
103 },
104 example! {
105 title: "Encrypt IPv6 address with AES128",
106 source: r#"encrypt_ip!("2001:db8::1", "sixteen byte key", "aes128")"#,
107 result: Ok("c0e6:eb35:6887:f554:4c65:8ace:17ca:6c6a"),
108 },
109 example! {
110 title: "Encrypt IPv4 address with prefix-preserving mode",
111 source: r#"encrypt_ip!("192.168.1.1", "thirty-two bytes key for pfx use", "pfx")"#,
112 result: Ok("33.245.248.61"),
113 },
114 example! {
115 title: "Encrypt IPv6 address with prefix-preserving mode",
116 source: r#"encrypt_ip!("2001:db8::1", "thirty-two bytes key for ipv6pfx", "pfx")"#,
117 result: Ok("88bd:d2bf:8865:8c4d:84b:44f6:6077:72c9"),
118 },
119 ]
120 }
121
122 fn compile(
123 &self,
124 _state: &state::TypeState,
125 _ctx: &mut FunctionCompileContext,
126 arguments: ArgumentList,
127 ) -> Compiled {
128 let ip = arguments.required("ip");
129 let key = arguments.required("key");
130 let mode = arguments.required("mode");
131
132 Ok(EncryptIpFn { ip, key, mode }.as_expr())
133 }
134}
135
136#[derive(Debug, Clone)]
137struct EncryptIpFn {
138 ip: Box<dyn Expression>,
139 key: Box<dyn Expression>,
140 mode: Box<dyn Expression>,
141}
142
143impl FunctionExpression for EncryptIpFn {
144 fn resolve(&self, ctx: &mut Context) -> Resolved {
145 let ip = self.ip.resolve(ctx)?;
146 let key = self.key.resolve(ctx)?;
147 let mode = self.mode.resolve(ctx)?;
148 encrypt_ip(&ip, key, &mode)
149 }
150
151 fn type_def(&self, _: &TypeState) -> TypeDef {
152 TypeDef::bytes().fallible()
153 }
154}
155
156#[cfg(test)]
157mod tests {
158 use super::*;
159 use crate::value;
160
161 test_function![
162 encrypt_ip => EncryptIp;
163
164 ipv4_aes128 {
165 args: func_args![
166 ip: "192.168.1.1",
167 key: value!(b"\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0c\x0d\x0e\x0f\x10"),
168 mode: "aes128"
169 ],
170 want: Ok(value!("a6d8:a149:6bcf:b175:bad6:3e56:d72d:4fdb")),
171 tdef: TypeDef::bytes().fallible(),
172 }
173
174 ipv4_pfx {
175 args: func_args![
176 ip: "192.168.1.1",
177 key: value!(b"\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0c\x0d\x0e\x0f\x10\x11\x12\x13\x14\x15\x16\x17\x18\x19\x1a\x1b\x1c\x1d\x1e\x1f\x20"),
178 mode: "pfx"
179 ],
180 want: Ok(value!("194.20.195.96")),
181 tdef: TypeDef::bytes().fallible(),
182 }
183
184 invalid_mode {
185 args: func_args![
186 ip: "192.168.1.1",
187 key: value!(b"\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0c\x0d\x0e\x0f\x10"),
188 mode: "invalid"
189 ],
190 want: Err("Invalid mode 'invalid'. Must be 'aes128' or 'pfx'"),
191 tdef: TypeDef::bytes().fallible(),
192 }
193
194 invalid_ip {
195 args: func_args![
196 ip: "not an ip",
197 key: value!(b"\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0c\x0d\x0e\x0f\x10"),
198 mode: "aes128"
199 ],
200 want: Err("unable to parse IP address: invalid IP address syntax"),
201 tdef: TypeDef::bytes().fallible(),
202 }
203
204 invalid_key_size_ipv4_aes128 {
205 args: func_args![
206 ip: "192.168.1.1",
207 key: value!(b"short"),
208 mode: "aes128"
209 ],
210 want: Err("aes128 mode requires a 16-byte key for IPv4"),
211 tdef: TypeDef::bytes().fallible(),
212 }
213
214 invalid_key_size_ipv4_pfx {
215 args: func_args![
216 ip: "192.168.1.1",
217 key: value!(b"\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0c\x0d\x0e\x0f\x10"),
218 mode: "pfx"
219 ],
220 want: Err("pfx mode requires a 32-byte key for IPv4"),
221 tdef: TypeDef::bytes().fallible(),
222 }
223 ];
224}