diff --git a/examples/threads/threads.txs b/examples/threads/threads.txs index 7396bb40..96d5b7db 100644 --- a/examples/threads/threads.txs +++ b/examples/threads/threads.txs @@ -6,12 +6,12 @@ spawn_fn = thread["spawn"] join_fn = thread["join"] do task_a() - time.sleep(100) + time.sleep(200) say "Task A completed on worker thread." end do task_b() - time.sleep(50) + time.sleep(20) say "Task B completed on worker thread." end diff --git a/stdlib/src/notification.rs b/stdlib/src/notification.rs index b4898fe7..1971cb28 100644 --- a/stdlib/src/notification.rs +++ b/stdlib/src/notification.rs @@ -1,3 +1,37 @@ +pub(crate) fn build_powershell_notification_args(title: &str, body: &str) -> Vec { + use base64::Engine; + + let b64_title = base64::engine::general_purpose::STANDARD.encode( + title + .encode_utf16() + .flat_map(|c| c.to_le_bytes()) + .collect::>(), + ); + let b64_body = base64::engine::general_purpose::STANDARD.encode( + body.encode_utf16() + .flat_map(|c| c.to_le_bytes()) + .collect::>(), + ); + + let script = format!( + "Add-Type -AssemblyName PresentationFramework; [System.Windows.MessageBox]::Show([System.Text.Encoding]::Unicode.GetString([System.Convert]::FromBase64String('{}')), [System.Text.Encoding]::Unicode.GetString([System.Convert]::FromBase64String('{}')))", + b64_body, b64_title + ); + let b64_script = base64::engine::general_purpose::STANDARD.encode( + script + .encode_utf16() + .flat_map(|c| c.to_le_bytes()) + .collect::>(), + ); + + vec![ + "-NoProfile".to_string(), + "-NonInteractive".to_string(), + "-EncodedCommand".to_string(), + b64_script, + ] +} + use crate::{StdFunction, StdlibModule, StdlibRegistry}; use std::collections::HashMap; use std::rc::Rc; @@ -32,21 +66,10 @@ impl StdlibRegistry { let body = args[1].to_string(); #[cfg(target_os = "windows")] { - use base64::Engine; use std::process::Command; - let b64_title = base64::engine::general_purpose::STANDARD.encode(title.encode_utf16().flat_map(|c| c.to_le_bytes()).collect::>()); - let b64_body = base64::engine::general_purpose::STANDARD.encode(body.encode_utf16().flat_map(|c| c.to_le_bytes()).collect::>()); - - let script = format!( - "Add-Type -AssemblyName PresentationFramework; [System.Windows.MessageBox]::Show([System.Text.Encoding]::Unicode.GetString([System.Convert]::FromBase64String('{}')), [System.Text.Encoding]::Unicode.GetString([System.Convert]::FromBase64String('{}')))", - b64_body, b64_title - ); - let b64_script = base64::engine::general_purpose::STANDARD.encode(script.encode_utf16().flat_map(|c| c.to_le_bytes()).collect::>()); - - let _ = Command::new("powershell") - .args(["-EncodedCommand", &b64_script]) - .spawn(); + let args = build_powershell_notification_args(&title, &body); + let _ = Command::new("powershell").args(&args).spawn(); } #[cfg(not(target_os = "windows"))] { @@ -127,4 +150,31 @@ mod tests { assert_eq!(alert.name(), "alert"); assert_eq!(alert.arity(), 1); } + + #[test] + fn test_powershell_notification_args_security() { + let title = "Test Title'; Remove-Item -Path C:\\ -Recurse -Force; '"; + let body = "Test Body\n\"'$$\x00"; + + let args = build_powershell_notification_args(title, body); + + assert_eq!(args[0], "-NoProfile"); + assert_eq!(args[1], "-NonInteractive"); + assert_eq!(args[2], "-EncodedCommand"); + + use base64::Engine; + let script_bytes = base64::engine::general_purpose::STANDARD + .decode(&args[3]) + .unwrap(); + let u16_chars: Vec = script_bytes + .chunks_exact(2) + .map(|chunk| u16::from_le_bytes([chunk[0], chunk[1]])) + .collect(); + let script = String::from_utf16(&u16_chars).unwrap(); + + assert!(script.contains("Add-Type -AssemblyName PresentationFramework")); + assert!(script.contains("[System.Windows.MessageBox]::Show")); + + assert!(!script.contains("Remove-Item")); + } }