diff --git a/crates/ironposh-client-core/src/pipeline.rs b/crates/ironposh-client-core/src/pipeline.rs index 4cb4999..ab352e5 100644 --- a/crates/ironposh-client-core/src/pipeline.rs +++ b/crates/ironposh-client-core/src/pipeline.rs @@ -73,7 +73,7 @@ pub struct ExecutionResult { /// This is owned and managed by the `RunspacePool`. #[derive(Debug, Clone)] pub struct Pipeline { - pub(crate) state: PsInvocationState, + state: PsInvocationState, pub(crate) commands: Vec, pub(crate) results: ExecutionResult, } @@ -98,6 +98,25 @@ impl Pipeline { pub(crate) fn add_command(&mut self, command: PipelineCommand) { self.commands.push(command); } + + /// Returns the current invocation state of the pipeline. + pub(crate) fn state(&self) -> PsInvocationState { + self.state + } + + /// Sets the invocation state of the pipeline. + pub(crate) fn set_state(&mut self, state: PsInvocationState) { + self.state = state; + } + + /// Returns `true` when the pipeline has reached a terminal state + /// (`Completed`, `Failed`, or `Stopped`). + pub(crate) fn is_terminal(&self) -> bool { + matches!( + self.state, + PsInvocationState::Completed | PsInvocationState::Failed | PsInvocationState::Stopped + ) + } } impl Pipeline { diff --git a/crates/ironposh-client-core/src/runspace_pool/crypto.rs b/crates/ironposh-client-core/src/runspace_pool/crypto.rs new file mode 100644 index 0000000..07ab9b8 --- /dev/null +++ b/crates/ironposh-client-core/src/runspace_pool/crypto.rs @@ -0,0 +1,120 @@ +//! PSRP session crypto: SecureString encryption and the key-exchange state. +//! +//! MS-PSRP encrypts `SecureString` values with a session key negotiated via +//! `PUBLIC_KEY` / `ENCRYPTED_SESSION_KEY` (AES-256-CBC, zero IV). This module +//! holds that state and the in-place encryption walk; the runspace pool drives +//! the key exchange and calls in here when serializing values. + +use aes::Aes256; +use cipher::block_padding::Pkcs7; +use cipher::{BlockModeEncrypt, KeyIvInit}; +use tracing::debug; + +#[derive(Debug)] +pub(super) struct KeyExchangeState { + pub(super) private_key: rsa::RsaPrivateKey, + pub(super) session_key: Option>, +} + +pub(super) fn encrypt_secure_strings_in_value_rec( + value: &mut ironposh_psrp::PsValue, + session_key: Option<&[u8]>, +) -> Result<(), crate::PwshCoreError> { + use ironposh_psrp::{ComplexObjectContent, Container, PsPrimitiveValue, PsValue}; + + match value { + PsValue::Primitive(PsPrimitiveValue::SecureString(bytes)) => { + let Some(session_key) = session_key else { + return Err(crate::PwshCoreError::InvalidResponse( + "SecureString encountered but PSRP session key is not established".into(), + )); + }; + encrypt_secure_string_bytes_in_place(bytes, session_key)?; + } + PsValue::Primitive(_) => {} + PsValue::Object(obj) => { + for value in obj.properties.values_mut() { + encrypt_secure_strings_in_value_rec(value, session_key)?; + } + + match &mut obj.content { + ComplexObjectContent::ExtendedPrimitive(p) => { + if let PsPrimitiveValue::SecureString(bytes) = p { + let Some(session_key) = session_key else { + return Err(crate::PwshCoreError::InvalidResponse( + "SecureString encountered but PSRP session key is not established" + .into(), + )); + }; + encrypt_secure_string_bytes_in_place(bytes, session_key)?; + } + } + ComplexObjectContent::Container( + Container::Stack(items) | Container::Queue(items) | Container::List(items), + ) => { + for item in items.iter_mut() { + encrypt_secure_strings_in_value_rec(item, session_key)?; + } + } + ComplexObjectContent::Container(Container::Dictionary(dict)) => { + for (_k, v) in dict.iter_mut() { + encrypt_secure_strings_in_value_rec(v, session_key)?; + } + } + ComplexObjectContent::Standard | ComplexObjectContent::PsEnums(_) => {} + } + } + } + + Ok(()) +} + +fn encrypt_secure_string_bytes_in_place( + bytes: &mut Vec, + session_key: &[u8], +) -> Result<(), crate::PwshCoreError> { + if session_key.len() != 32 { + return Err(crate::PwshCoreError::InvalidResponse( + format!( + "PSRP SecureString encryption requires 32-byte session key; got {}", + session_key.len() + ) + .into(), + )); + } + + // PowerShell's PSRP SecureString encryption uses AES-256-CBC with a zero IV. + // The payload is the ciphertext bytes only (base64 encoded). + let iv = [0u8; 16]; + + let encryptor = cbc::Encryptor::::new_from_slices(session_key, &iv).map_err(|e| { + crate::PwshCoreError::InvalidResponse( + format!("Failed to initialize AES encryptor: {e}").into(), + ) + })?; + + // MS-PSRP SecureString payload is UTF-16LE plaintext encrypted with AES-256-CBC. + let msg_len = bytes.len(); + let pad = 16 - (msg_len % 16); + let mut buf = bytes.clone(); + buf.resize(msg_len + pad, 0); + let ciphertext = encryptor + .encrypt_padded::(&mut buf, msg_len) + .map_err(|e| { + crate::PwshCoreError::InvalidResponse( + format!("Failed to encrypt SecureString (padding): {e}").into(), + ) + })?; + + let out = ciphertext.to_vec(); + + debug!( + session_key_len = session_key.len(), + plaintext_len = msg_len, + encrypted_len = out.len(), + "encrypted SecureString payload" + ); + + *bytes = out; + Ok(()) +} diff --git a/crates/ironposh-client-core/src/runspace_pool/host_call.rs b/crates/ironposh-client-core/src/runspace_pool/host_call.rs new file mode 100644 index 0000000..44d139a --- /dev/null +++ b/crates/ironposh-client-core/src/runspace_pool/host_call.rs @@ -0,0 +1,75 @@ +//! Pure host-call construction and classification helpers extracted from +//! `runspace_pool::pool`. +//! +//! These functions are behavior-preserving and free of any `RunspacePool` +//! state: they only parse a `PsValue` into a [`HostCall`] and classify an +//! existing [`HostCall`]. Anything that needs pool internals (fragmenter, id, +//! shell, the `pending_host_calls` queue) stays in `pool.rs`. + +use tracing::debug; +use uuid::Uuid; + +use crate::{ + PwshCoreError, + host::{HostCall, HostCallScope}, +}; +use ironposh_psrp::PsValue; + +/// Parse a `PipelineHostCall` `PsValue` into a [`HostCall`] scoped to a pipeline. +/// +/// This is the pure parsing body previously inlined in +/// `RunspacePool::handle_pipeline_host_call`; it does not touch pool state. +pub(super) fn pipeline_host_call_from( + ps_value: PsValue, + stream_name: &str, + command_id: Option<&Uuid>, +) -> Result { + let PsValue::Object(pipeline_host_call) = ps_value else { + return Err(PwshCoreError::InvalidResponse( + "Expected PipelineHostCall as PsValue::Object".into(), + )); + }; + + let pipeline_host_call = ironposh_psrp::PipelineHostCall::try_from(pipeline_host_call)?; + + debug!( + ?pipeline_host_call, + stream_name = stream_name, + command_id = ?command_id, + method = ?pipeline_host_call.method, + parameters = ?pipeline_host_call.parameters, + "Received PipelineHostCall" + ); + + // Question: Can we have a Optional command id here? + let Some(command_id) = command_id else { + return Err(PwshCoreError::InvalidResponse( + "Expected command_id to be Some".into(), + )); + }; + + let scope = HostCallScope::Pipeline { + command_id: command_id.to_owned(), + }; + + HostCall::try_from_pipeline(scope, pipeline_host_call).map_err(|e| { + PwshCoreError::InvalidResponse(format!("Failed to parse host call: {e}").into()) + }) +} + +/// Classify whether a host call requires a PSRP session key to be established +/// before it can be answered (because it transports secure-string data). +pub(super) fn needs_session_key(host_call: &HostCall) -> bool { + match host_call { + HostCall::ReadLineAsSecureString { .. } + | HostCall::PromptForCredential1 { .. } + | HostCall::PromptForCredential2 { .. } => true, + HostCall::Prompt { transport } => { + let (_, _, fields) = &transport.params; + fields + .iter() + .any(|f| f.parameter_type.contains("SecureString")) + } + _ => false, + } +} diff --git a/crates/ironposh-client-core/src/runspace_pool/mod.rs b/crates/ironposh-client-core/src/runspace_pool/mod.rs index 445f409..abf5543 100644 --- a/crates/ironposh-client-core/src/runspace_pool/mod.rs +++ b/crates/ironposh-client-core/src/runspace_pool/mod.rs @@ -1,7 +1,9 @@ pub mod creator; +mod crypto; pub mod enums; pub mod expect_shell_connected; pub mod expect_shell_created; +mod host_call; pub mod pool; pub mod types; diff --git a/crates/ironposh-client-core/src/runspace_pool/pool.rs b/crates/ironposh-client-core/src/runspace_pool/pool.rs index a8bd4c1..821a6b3 100644 --- a/crates/ironposh-client-core/src/runspace_pool/pool.rs +++ b/crates/ironposh-client-core/src/runspace_pool/pool.rs @@ -18,14 +18,11 @@ use rsa::traits::PublicKeyParts; use rsa::{RsaPrivateKey, pkcs1v15::Pkcs1v15Encrypt}; use tracing::{debug, error, info, instrument, trace, warn}; -use aes::Aes256; -use cipher::block_padding::Pkcs7; -use cipher::{BlockModeEncrypt, KeyIvInit}; use uuid::Uuid; use crate::{ PwshCoreError, - host::{HostCall, HostCallScope}, + host::HostCall, pipeline::{Pipeline, PipelineCommand, PipelineSpec}, powershell::PipelineHandle, runspace::win_rs::WinRunspace, @@ -116,12 +113,6 @@ pub enum AcceptResponsResult { }, } -#[derive(Debug)] -pub(super) struct KeyExchangeState { - private_key: rsa::RsaPrivateKey, - session_key: Option>, -} - #[derive(Debug)] pub struct RunspacePool { pub(super) id: uuid::Uuid, @@ -140,114 +131,11 @@ pub struct RunspacePool { pub(super) pipelines: HashMap, pub(super) fragmenter: fragmentation::Fragmenter, pub(super) desired_stream_is_pooling: bool, - pub(super) key_exchange: Option, + pub(super) key_exchange: Option, pub(super) psrp_key_exchange_pending: bool, pub(super) pending_host_calls: VecDeque, } -fn encrypt_secure_strings_in_value_rec( - value: &mut ironposh_psrp::PsValue, - session_key: Option<&[u8]>, -) -> Result<(), crate::PwshCoreError> { - use ironposh_psrp::{ComplexObjectContent, Container, PsPrimitiveValue, PsValue}; - - match value { - PsValue::Primitive(PsPrimitiveValue::SecureString(bytes)) => { - let Some(session_key) = session_key else { - return Err(crate::PwshCoreError::InvalidResponse( - "SecureString encountered but PSRP session key is not established".into(), - )); - }; - encrypt_secure_string_bytes_in_place(bytes, session_key)?; - } - PsValue::Primitive(_) => {} - PsValue::Object(obj) => { - for value in obj.properties.values_mut() { - encrypt_secure_strings_in_value_rec(value, session_key)?; - } - - match &mut obj.content { - ComplexObjectContent::ExtendedPrimitive(p) => { - if let PsPrimitiveValue::SecureString(bytes) = p { - let Some(session_key) = session_key else { - return Err(crate::PwshCoreError::InvalidResponse( - "SecureString encountered but PSRP session key is not established" - .into(), - )); - }; - encrypt_secure_string_bytes_in_place(bytes, session_key)?; - } - } - ComplexObjectContent::Container( - Container::Stack(items) | Container::Queue(items) | Container::List(items), - ) => { - for item in items.iter_mut() { - encrypt_secure_strings_in_value_rec(item, session_key)?; - } - } - ComplexObjectContent::Container(Container::Dictionary(dict)) => { - for (_k, v) in dict.iter_mut() { - encrypt_secure_strings_in_value_rec(v, session_key)?; - } - } - ComplexObjectContent::Standard | ComplexObjectContent::PsEnums(_) => {} - } - } - } - - Ok(()) -} - -fn encrypt_secure_string_bytes_in_place( - bytes: &mut Vec, - session_key: &[u8], -) -> Result<(), crate::PwshCoreError> { - if session_key.len() != 32 { - return Err(crate::PwshCoreError::InvalidResponse( - format!( - "PSRP SecureString encryption requires 32-byte session key; got {}", - session_key.len() - ) - .into(), - )); - } - - // PowerShell's PSRP SecureString encryption uses AES-256-CBC with a zero IV. - // The payload is the ciphertext bytes only (base64 encoded). - let iv = [0u8; 16]; - - let encryptor = cbc::Encryptor::::new_from_slices(session_key, &iv).map_err(|e| { - crate::PwshCoreError::InvalidResponse( - format!("Failed to initialize AES encryptor: {e}").into(), - ) - })?; - - // MS-PSRP SecureString payload is UTF-16LE plaintext encrypted with AES-256-CBC. - let msg_len = bytes.len(); - let pad = 16 - (msg_len % 16); - let mut buf = bytes.clone(); - buf.resize(msg_len + pad, 0); - let ciphertext = encryptor - .encrypt_padded::(&mut buf, msg_len) - .map_err(|e| { - crate::PwshCoreError::InvalidResponse( - format!("Failed to encrypt SecureString (padding): {e}").into(), - ) - })?; - - let out = ciphertext.to_vec(); - - debug!( - session_key_len = session_key.len(), - plaintext_len = msg_len, - encrypted_len = out.len(), - "encrypted SecureString payload" - ); - - *bytes = out; - Ok(()) -} - impl RunspacePool { pub fn encrypt_secure_strings_in_value( &self, @@ -257,7 +145,7 @@ impl RunspacePool { .key_exchange .as_ref() .and_then(|s| s.session_key.as_deref()); - encrypt_secure_strings_in_value_rec(value, session_key) + super::crypto::encrypt_secure_strings_in_value_rec(value, session_key) } /// Build the negotiation payload shared by [`Self::open`] and @@ -723,7 +611,7 @@ impl RunspacePool { "Pipeline not found for command response".into(), ) })? - .state = PsInvocationState::Running; + .set_state(PsInvocationState::Running); result.push(AcceptResponsResult::ReceiveResponse { desired_streams: vec![DesiredStream::stdout_for_command(pipeline_id)], @@ -1197,18 +1085,7 @@ impl RunspacePool { })?; debug!(target: "host_call", host_call = ?host_call, "successfully created host call"); - let needs_session_key = match &host_call { - HostCall::ReadLineAsSecureString { .. } - | HostCall::PromptForCredential1 { .. } - | HostCall::PromptForCredential2 { .. } => true, - HostCall::Prompt { transport } => { - let (_, _, fields) = &transport.params; - fields - .iter() - .any(|f| f.parameter_type.contains("SecureString")) - } - _ => false, - }; + let needs_session_key = super::host_call::needs_session_key(&host_call); let has_session_key = self .key_exchange @@ -1555,7 +1432,7 @@ impl RunspacePool { PwshCoreError::InvalidResponse("Pipeline not found for command_id".into()) })?; // Update the pipeline state - pipeline.state = PsInvocationState::from(pipeline_state.pipeline_state); + pipeline.set_state(PsInvocationState::from(pipeline_state.pipeline_state)); Ok(()) } @@ -1571,7 +1448,7 @@ impl RunspacePool { .ok_or(PwshCoreError::InvalidState("Pipeline handle not found"))?; // Set pipeline state to Running - pipeline.state = PsInvocationState::Running; + pipeline.set_state(PsInvocationState::Running); info!(pipeline_id = %handle.id(), "Invoking pipeline"); // Convert business pipeline to protocol pipeline and build CreatePipeline message @@ -1615,17 +1492,14 @@ impl RunspacePool { error!(pipeline_id = ?&handle.id(), "Pipeline handle not found "); })?; - if pipeline.state == PsInvocationState::Stopped - || pipeline.state == PsInvocationState::Completed - || pipeline.state == PsInvocationState::Failed - { + if pipeline.is_terminal() { return Err(PwshCoreError::InvalidState( "Cannot kill a pipeline that is already stopped, completed, or failed", )); } // Set pipeline state to Stopping - pipeline.state = PsInvocationState::Stopping; + pipeline.set_state(PsInvocationState::Stopping); info!(pipeline_id = %handle.id(), "Killing pipeline"); let request = self @@ -1642,37 +1516,7 @@ impl RunspacePool { stream_name: &str, command_id: Option<&Uuid>, ) -> Result { - let PsValue::Object(pipeline_host_call) = ps_value else { - return Err(PwshCoreError::InvalidResponse( - "Expected PipelineHostCall as PsValue::Object".into(), - )); - }; - - let pipeline_host_call = ironposh_psrp::PipelineHostCall::try_from(pipeline_host_call)?; - - debug!( - ?pipeline_host_call, - stream_name = stream_name, - command_id = ?command_id, - method = ?pipeline_host_call.method, - parameters = ?pipeline_host_call.parameters, - "Received PipelineHostCall" - ); - - // Question: Can we have a Optional command id here? - let Some(command_id) = command_id else { - return Err(PwshCoreError::InvalidResponse( - "Expected command_id to be Some".into(), - )); - }; - - let scope = HostCallScope::Pipeline { - command_id: command_id.to_owned(), - }; - - HostCall::try_from_pipeline(scope, pipeline_host_call).map_err(|e| { - crate::PwshCoreError::InvalidResponse(format!("Failed to parse host call: {e}").into()) - }) + super::host_call::pipeline_host_call_from(ps_value, stream_name, command_id) } /// Send a pipeline host response to the server @@ -1816,13 +1660,15 @@ impl RunspacePool { Ok(xml) } - fn ensure_key_exchange_state(&mut self) -> Result<&mut KeyExchangeState, PwshCoreError> { + fn ensure_key_exchange_state( + &mut self, + ) -> Result<&mut super::crypto::KeyExchangeState, PwshCoreError> { if self.key_exchange.is_none() { let mut rng = rand::thread_rng(); let private_key = RsaPrivateKey::new(&mut rng, 2048).map_err(|e| { PwshCoreError::InternalError(format!("failed to generate RSA keypair: {e}")) })?; - self.key_exchange = Some(KeyExchangeState { + self.key_exchange = Some(super::crypto::KeyExchangeState { private_key, session_key: None, }); @@ -1899,7 +1745,7 @@ impl RunspacePool { .get_mut(&powershell.id()) .ok_or(PwshCoreError::InvalidState("Pipeline handle not found"))?; - if pipeline.state != PsInvocationState::NotStarted { + if pipeline.state() != PsInvocationState::NotStarted { return Err(PwshCoreError::InvalidState( "Cannot add to a pipeline that has already been started", ));