1use std::collections::HashSet;
4
5use serde_json::Value;
6
7use crate::functions::Function;
8use crate::interpreter::SearchResult;
9use crate::registry::register_if_enabled;
10use crate::{Context, Runtime, arg, defn};
11
12use base64::{
13 Engine,
14 engine::general_purpose::{STANDARD as BASE64_STANDARD, URL_SAFE_NO_PAD as BASE64_URL_SAFE},
15};
16
17pub fn register_filtered(runtime: &mut Runtime, enabled: &HashSet<&str>) {
19 register_if_enabled(
20 runtime,
21 "base64_encode",
22 enabled,
23 Box::new(Base64EncodeFn::new()),
24 );
25 register_if_enabled(
26 runtime,
27 "base64_decode",
28 enabled,
29 Box::new(Base64DecodeFn::new()),
30 );
31 register_if_enabled(
32 runtime,
33 "base64url_decode",
34 enabled,
35 Box::new(Base64UrlDecodeFn::new()),
36 );
37 register_if_enabled(
38 runtime,
39 "base64url_encode",
40 enabled,
41 Box::new(Base64UrlEncodeFn::new()),
42 );
43 register_if_enabled(runtime, "hex_encode", enabled, Box::new(HexEncodeFn::new()));
44 register_if_enabled(runtime, "hex_decode", enabled, Box::new(HexDecodeFn::new()));
45 register_if_enabled(runtime, "jwt_decode", enabled, Box::new(JwtDecodeFn::new()));
46 register_if_enabled(runtime, "jwt_header", enabled, Box::new(JwtHeaderFn::new()));
47 register_if_enabled(
48 runtime,
49 "html_escape",
50 enabled,
51 Box::new(HtmlEscapeFn::new()),
52 );
53 register_if_enabled(
54 runtime,
55 "html_unescape",
56 enabled,
57 Box::new(HtmlUnescapeFn::new()),
58 );
59 register_if_enabled(
60 runtime,
61 "shell_escape",
62 enabled,
63 Box::new(ShellEscapeFn::new()),
64 );
65}
66
67defn!(Base64EncodeFn, vec![arg!(string)], None);
72
73impl Function for Base64EncodeFn {
74 fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
75 self.signature.validate(args, ctx)?;
76
77 let input = args[0].as_str().ok_or_else(|| {
78 crate::JmespathError::from_ctx(
79 ctx,
80 crate::ErrorReason::Parse("Expected string argument".to_owned()),
81 )
82 })?;
83
84 let encoded = BASE64_STANDARD.encode(input.as_bytes());
85 Ok(Value::String(encoded))
86 }
87}
88
89defn!(Base64DecodeFn, vec![arg!(string)], None);
94
95impl Function for Base64DecodeFn {
96 fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
97 self.signature.validate(args, ctx)?;
98
99 let input = args[0].as_str().ok_or_else(|| {
100 crate::JmespathError::from_ctx(
101 ctx,
102 crate::ErrorReason::Parse("Expected string argument".to_owned()),
103 )
104 })?;
105
106 match BASE64_STANDARD.decode(input.as_bytes()) {
107 Ok(decoded) => {
108 let s = String::from_utf8(decoded).map_err(|_| {
109 crate::JmespathError::from_ctx(
110 ctx,
111 crate::ErrorReason::Parse("Decoded bytes are not valid UTF-8".to_owned()),
112 )
113 })?;
114 Ok(Value::String(s))
115 }
116 Err(_) => Err(crate::JmespathError::from_ctx(
117 ctx,
118 crate::ErrorReason::Parse("Invalid base64 input".to_owned()),
119 )),
120 }
121 }
122}
123
124defn!(Base64UrlEncodeFn, vec![arg!(string)], None);
129
130impl Function for Base64UrlEncodeFn {
131 fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
132 self.signature.validate(args, ctx)?;
133
134 let input = args[0].as_str().ok_or_else(|| {
135 crate::JmespathError::from_ctx(
136 ctx,
137 crate::ErrorReason::Parse("Expected string argument".to_owned()),
138 )
139 })?;
140
141 let encoded = BASE64_URL_SAFE.encode(input.as_bytes());
142 Ok(Value::String(encoded))
143 }
144}
145
146defn!(Base64UrlDecodeFn, vec![arg!(string)], None);
151
152impl Function for Base64UrlDecodeFn {
153 fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
154 self.signature.validate(args, ctx)?;
155
156 let input = args[0].as_str().ok_or_else(|| {
157 crate::JmespathError::from_ctx(
158 ctx,
159 crate::ErrorReason::Parse("Expected string argument".to_owned()),
160 )
161 })?;
162
163 match BASE64_URL_SAFE.decode(input.as_bytes()) {
164 Ok(decoded) => {
165 let s = String::from_utf8(decoded).map_err(|_| {
166 crate::JmespathError::from_ctx(
167 ctx,
168 crate::ErrorReason::Parse("Decoded bytes are not valid UTF-8".to_owned()),
169 )
170 })?;
171 Ok(Value::String(s))
172 }
173 Err(_) => Err(crate::JmespathError::from_ctx(
174 ctx,
175 crate::ErrorReason::Parse("Invalid base64url input".to_owned()),
176 )),
177 }
178 }
179}
180
181defn!(HexEncodeFn, vec![arg!(string)], None);
186
187impl Function for HexEncodeFn {
188 fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
189 self.signature.validate(args, ctx)?;
190
191 let input = args[0].as_str().ok_or_else(|| {
192 crate::JmespathError::from_ctx(
193 ctx,
194 crate::ErrorReason::Parse("Expected string argument".to_owned()),
195 )
196 })?;
197
198 let encoded = hex::encode(input.as_bytes());
199 Ok(Value::String(encoded))
200 }
201}
202
203defn!(HexDecodeFn, vec![arg!(string)], None);
208
209impl Function for HexDecodeFn {
210 fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
211 self.signature.validate(args, ctx)?;
212
213 let input = args[0].as_str().ok_or_else(|| {
214 crate::JmespathError::from_ctx(
215 ctx,
216 crate::ErrorReason::Parse("Expected string argument".to_owned()),
217 )
218 })?;
219
220 match hex::decode(input) {
221 Ok(decoded) => {
222 match String::from_utf8(decoded) {
224 Ok(s) => Ok(Value::String(s)),
225 Err(_) => Ok(Value::Null),
226 }
227 }
228 Err(_) => Ok(Value::Null),
230 }
231 }
232}
233
234fn decode_jwt_part(part: &str) -> Option<serde_json::Value> {
240 let decoded = BASE64_URL_SAFE.decode(part).ok()?;
242 let json_str = String::from_utf8(decoded).ok()?;
243 serde_json::from_str(&json_str).ok()
244}
245
246defn!(JwtDecodeFn, vec![arg!(string)], None);
251
252impl Function for JwtDecodeFn {
253 fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
254 self.signature.validate(args, ctx)?;
255
256 let token = args[0].as_str().ok_or_else(|| {
257 crate::JmespathError::from_ctx(
258 ctx,
259 crate::ErrorReason::Parse("Expected string argument".to_owned()),
260 )
261 })?;
262
263 let parts: Vec<&str> = token.split('.').collect();
265 if parts.len() != 3 {
266 return Ok(Value::Null);
267 }
268
269 match decode_jwt_part(parts[1]) {
271 Some(json) => Ok(json),
272 None => Ok(Value::Null),
273 }
274 }
275}
276
277defn!(JwtHeaderFn, vec![arg!(string)], None);
282
283impl Function for JwtHeaderFn {
284 fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
285 self.signature.validate(args, ctx)?;
286
287 let token = args[0].as_str().ok_or_else(|| {
288 crate::JmespathError::from_ctx(
289 ctx,
290 crate::ErrorReason::Parse("Expected string argument".to_owned()),
291 )
292 })?;
293
294 let parts: Vec<&str> = token.split('.').collect();
296 if parts.len() != 3 {
297 return Ok(Value::Null);
298 }
299
300 match decode_jwt_part(parts[0]) {
302 Some(json) => Ok(json),
303 None => Ok(Value::Null),
304 }
305 }
306}
307
308defn!(HtmlEscapeFn, vec![arg!(string)], None);
313
314impl Function for HtmlEscapeFn {
315 fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
316 self.signature.validate(args, ctx)?;
317
318 let s = args[0].as_str().ok_or_else(|| {
319 crate::JmespathError::from_ctx(
320 ctx,
321 crate::ErrorReason::Parse("Expected string argument".to_owned()),
322 )
323 })?;
324
325 let escaped = s
326 .replace('&', "&")
327 .replace('<', "<")
328 .replace('>', ">")
329 .replace('"', """)
330 .replace('\'', "'");
331
332 Ok(Value::String(escaped))
333 }
334}
335
336defn!(HtmlUnescapeFn, vec![arg!(string)], None);
341
342impl Function for HtmlUnescapeFn {
343 fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
344 self.signature.validate(args, ctx)?;
345
346 let s = args[0].as_str().ok_or_else(|| {
347 crate::JmespathError::from_ctx(
348 ctx,
349 crate::ErrorReason::Parse("Expected string argument".to_owned()),
350 )
351 })?;
352
353 let unescaped = s
355 .replace("'", "'")
356 .replace("'", "'")
357 .replace("'", "'")
358 .replace(""", "\"")
359 .replace(">", ">")
360 .replace("<", "<")
361 .replace("&", "&");
362
363 Ok(Value::String(unescaped))
364 }
365}
366
367defn!(ShellEscapeFn, vec![arg!(string)], None);
372
373impl Function for ShellEscapeFn {
374 fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
375 self.signature.validate(args, ctx)?;
376
377 let s = args[0].as_str().ok_or_else(|| {
378 crate::JmespathError::from_ctx(
379 ctx,
380 crate::ErrorReason::Parse("Expected string argument".to_owned()),
381 )
382 })?;
383
384 let escaped = format!("'{}'", s.replace('\'', "'\\''"));
387
388 Ok(Value::String(escaped))
389 }
390}
391
392#[cfg(test)]
393mod tests {
394 use crate::Runtime;
395 use serde_json::json;
396
397 fn setup_runtime() -> Runtime {
398 Runtime::builder()
399 .with_standard()
400 .with_all_extensions()
401 .build()
402 }
403
404 #[test]
405 fn test_base64_encode() {
406 let runtime = setup_runtime();
407 let expr = runtime.compile("base64_encode(@)").unwrap();
408 let data = json!("hello");
409 let result = expr.search(&data).unwrap();
410 assert_eq!(result, json!("aGVsbG8="));
411 }
412
413 #[test]
414 fn test_base64_decode() {
415 let runtime = setup_runtime();
416 let expr = runtime.compile("base64_decode(@)").unwrap();
417 let data = json!("aGVsbG8=");
418 let result = expr.search(&data).unwrap();
419 assert_eq!(result, json!("hello"));
420 }
421
422 #[test]
423 fn test_hex_encode() {
424 let runtime = setup_runtime();
425 let expr = runtime.compile("hex_encode(@)").unwrap();
426 let data = json!("hello");
427 let result = expr.search(&data).unwrap();
428 assert_eq!(result, json!("68656c6c6f"));
429 }
430
431 #[test]
432 fn test_hex_decode() {
433 let runtime = setup_runtime();
434 let expr = runtime.compile("hex_decode(@)").unwrap();
435 let data = json!("68656c6c6f");
436 let result = expr.search(&data).unwrap();
437 assert_eq!(result, json!("hello"));
438 }
439
440 #[test]
441 fn test_hex_decode_invalid_returns_null() {
442 let runtime = setup_runtime();
443 let expr = runtime.compile("hex_decode(@)").unwrap();
444 let data = json!("invalid");
445 let result = expr.search(&data).unwrap();
446 assert_eq!(result, json!(null));
447 }
448
449 #[test]
450 fn test_hex_decode_odd_length_returns_null() {
451 let runtime = setup_runtime();
452 let expr = runtime.compile("hex_decode(@)").unwrap();
453 let data = json!("123");
454 let result = expr.search(&data).unwrap();
455 assert_eq!(result, json!(null));
456 }
457
458 const TEST_JWT: &str = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c";
464
465 #[test]
466 fn test_jwt_decode_payload() {
467 let runtime = setup_runtime();
468 let expr = runtime.compile("jwt_decode(@)").unwrap();
469 let data = json!(TEST_JWT);
470 let result = expr.search(&data).unwrap();
471
472 assert_eq!(result["sub"], json!("1234567890"));
474 assert_eq!(result["name"], json!("John Doe"));
475 assert_eq!(result["iat"], json!(1516239022));
476 }
477
478 #[test]
479 fn test_jwt_decode_extract_claim() {
480 let runtime = setup_runtime();
481 let expr = runtime.compile("jwt_decode(@).sub").unwrap();
482 let data = json!(TEST_JWT);
483 let result = expr.search(&data).unwrap();
484 assert_eq!(result, json!("1234567890"));
485 }
486
487 #[test]
488 fn test_jwt_header() {
489 let runtime = setup_runtime();
490 let expr = runtime.compile("jwt_header(@)").unwrap();
491 let data = json!(TEST_JWT);
492 let result = expr.search(&data).unwrap();
493
494 assert_eq!(result["alg"], json!("HS256"));
496 assert_eq!(result["typ"], json!("JWT"));
497 }
498
499 #[test]
500 fn test_jwt_header_extract_alg() {
501 let runtime = setup_runtime();
502 let expr = runtime.compile("jwt_header(@).alg").unwrap();
503 let data = json!(TEST_JWT);
504 let result = expr.search(&data).unwrap();
505 assert_eq!(result, json!("HS256"));
506 }
507
508 #[test]
509 fn test_jwt_decode_invalid_format() {
510 let runtime = setup_runtime();
511 let expr = runtime.compile("jwt_decode(@)").unwrap();
512
513 let data = json!("not-a-jwt");
515 let result = expr.search(&data).unwrap();
516 assert_eq!(result, json!(null));
517
518 let data = json!("part1.part2");
520 let result = expr.search(&data).unwrap();
521 assert_eq!(result, json!(null));
522 }
523
524 #[test]
525 fn test_jwt_decode_invalid_base64() {
526 let runtime = setup_runtime();
527 let expr = runtime.compile("jwt_decode(@)").unwrap();
528
529 let data = json!("!!!.@@@.###");
531 let result = expr.search(&data).unwrap();
532 assert_eq!(result, json!(null));
533 }
534
535 #[test]
536 fn test_jwt_decode_invalid_json() {
537 let runtime = setup_runtime();
538 let expr = runtime.compile("jwt_decode(@)").unwrap();
539
540 let data = json!("eyJhbGciOiJIUzI1NiJ9.bm90IGpzb24.sig");
542 let result = expr.search(&data).unwrap();
543 assert_eq!(result, json!(null));
544 }
545
546 #[test]
547 fn test_html_escape_basic() {
548 let runtime = setup_runtime();
549 let expr = runtime.compile("html_escape(@)").unwrap();
550 let data = json!("<div class=\"test\">Hello & goodbye</div>");
551 let result = expr.search(&data).unwrap();
552 assert_eq!(
553 result,
554 json!("<div class="test">Hello & goodbye</div>")
555 );
556 }
557
558 #[test]
559 fn test_html_escape_quotes() {
560 let runtime = setup_runtime();
561 let expr = runtime.compile("html_escape(@)").unwrap();
562 let data = json!("It's a \"test\"");
563 let result = expr.search(&data).unwrap();
564 assert_eq!(result, json!("It's a "test""));
565 }
566
567 #[test]
568 fn test_html_escape_no_change() {
569 let runtime = setup_runtime();
570 let expr = runtime.compile("html_escape(@)").unwrap();
571 let data = json!("Hello World");
572 let result = expr.search(&data).unwrap();
573 assert_eq!(result, json!("Hello World"));
574 }
575
576 #[test]
577 fn test_html_unescape_basic() {
578 let runtime = setup_runtime();
579 let expr = runtime.compile("html_unescape(@)").unwrap();
580 let data = json!("<div class="test">Hello & goodbye</div>");
581 let result = expr.search(&data).unwrap();
582 assert_eq!(result, json!("<div class=\"test\">Hello & goodbye</div>"));
583 }
584
585 #[test]
586 fn test_html_unescape_quotes() {
587 let runtime = setup_runtime();
588 let expr = runtime.compile("html_unescape(@)").unwrap();
589 let data = json!("It's a "test"");
590 let result = expr.search(&data).unwrap();
591 assert_eq!(result, json!("It's a \"test\""));
592 }
593
594 #[test]
595 fn test_html_roundtrip() {
596 let runtime = setup_runtime();
597 let escape = runtime.compile("html_escape(@)").unwrap();
598 let unescape = runtime.compile("html_unescape(@)").unwrap();
599 let original = "<script>alert('xss')</script>";
600 let data = json!(original);
601 let escaped = escape.search(&data).unwrap();
602 let roundtrip = unescape.search(&escaped).unwrap();
603 assert_eq!(roundtrip, json!(original));
604 }
605
606 #[test]
607 fn test_shell_escape_simple() {
608 let runtime = setup_runtime();
609 let expr = runtime.compile("shell_escape(@)").unwrap();
610 let data = json!("hello world");
611 let result = expr.search(&data).unwrap();
612 assert_eq!(result, json!("'hello world'"));
613 }
614
615 #[test]
616 fn test_shell_escape_with_single_quote() {
617 let runtime = setup_runtime();
618 let expr = runtime.compile("shell_escape(@)").unwrap();
619 let data = json!("it's here");
620 let result = expr.search(&data).unwrap();
621 assert_eq!(result, json!("'it'\\''s here'"));
622 }
623
624 #[test]
625 fn test_shell_escape_special_chars() {
626 let runtime = setup_runtime();
627 let expr = runtime.compile("shell_escape(@)").unwrap();
628 let data = json!("$HOME; rm -rf /");
629 let result = expr.search(&data).unwrap();
630 assert_eq!(result, json!("'$HOME; rm -rf /'"));
632 }
633
634 #[test]
635 fn test_shell_escape_empty() {
636 let runtime = setup_runtime();
637 let expr = runtime.compile("shell_escape(@)").unwrap();
638 let data = json!("");
639 let result = expr.search(&data).unwrap();
640 assert_eq!(result, json!("''"));
641 }
642
643 #[test]
644 fn test_shell_escape_multiple_quotes() {
645 let runtime = setup_runtime();
646 let expr = runtime.compile("shell_escape(@)").unwrap();
647 let data = json!("don't say 'hello'");
648 let result = expr.search(&data).unwrap();
649 assert_eq!(result, json!("'don'\\''t say '\\''hello'\\'''"));
650 }
651
652 #[test]
657 fn test_base64url_encode() {
658 let runtime = setup_runtime();
659 let expr = runtime.compile("base64url_encode(@)").unwrap();
660 let data = json!("hello");
661 let result = expr.search(&data).unwrap();
662 assert_eq!(result, json!("aGVsbG8"));
663 }
664
665 #[test]
666 fn test_base64url_decode() {
667 let runtime = setup_runtime();
668 let expr = runtime.compile("base64url_decode(@)").unwrap();
669 let data = json!("aGVsbG8");
670 let result = expr.search(&data).unwrap();
671 assert_eq!(result, json!("hello"));
672 }
673
674 #[test]
675 fn test_base64url_roundtrip() {
676 let runtime = setup_runtime();
677 let encode = runtime.compile("base64url_encode(@)").unwrap();
678 let decode = runtime.compile("base64url_decode(@)").unwrap();
679 let original = "hello world! 🌍";
680 let data = json!(original);
681 let encoded = encode.search(&data).unwrap();
682 let roundtrip = decode.search(&encoded).unwrap();
683 assert_eq!(roundtrip, json!(original));
684 }
685
686 #[test]
687 fn test_base64url_no_padding() {
688 let runtime = setup_runtime();
689 let expr = runtime.compile("base64url_encode(@)").unwrap();
690 let data = json!("test");
692 let result = expr.search(&data).unwrap();
693 let s = result.as_str().unwrap();
694 assert!(
695 !s.contains('='),
696 "base64url output should not contain padding"
697 );
698 assert_eq!(s, "dGVzdA");
699 }
700
701 #[test]
702 fn test_base64url_uses_url_safe_chars() {
703 let runtime = setup_runtime();
704 let encode_url = runtime.compile("base64url_encode(@)").unwrap();
705 let encode_std = runtime.compile("base64_encode(@)").unwrap();
706 let data = json!("subjects?_d");
711 let std_result = encode_std.search(&data).unwrap();
712 let url_result = encode_url.search(&data).unwrap();
713 let std_s = std_result.as_str().unwrap();
714 let url_s = url_result.as_str().unwrap();
715 assert!(!url_s.contains('+'), "base64url should not contain '+'");
717 assert!(!url_s.contains('/'), "base64url should not contain '/'");
718 if std_s.contains('+') || std_s.contains('/') {
720 assert_ne!(std_s.trim_end_matches('='), url_s);
721 }
722 }
723
724 #[test]
725 fn test_base64url_decode_invalid() {
726 let runtime = setup_runtime();
727 let expr = runtime.compile("base64url_decode(@)").unwrap();
728 let data = json!("!!!invalid!!!");
729 let result = expr.search(&data);
730 assert!(result.is_err());
731 }
732
733 #[test]
734 fn test_base64url_encode_empty() {
735 let runtime = setup_runtime();
736 let expr = runtime.compile("base64url_encode(@)").unwrap();
737 let data = json!("");
738 let result = expr.search(&data).unwrap();
739 assert_eq!(result, json!(""));
740 }
741
742 #[test]
743 fn test_base64url_decode_empty() {
744 let runtime = setup_runtime();
745 let expr = runtime.compile("base64url_decode(@)").unwrap();
746 let data = json!("");
747 let result = expr.search(&data).unwrap();
748 assert_eq!(result, json!(""));
749 }
750}