diff --git a/mitmproxy-rs/src/contentview.rs b/mitmproxy-rs/src/contentview.rs new file mode 100644 index 00000000..dac40b24 --- /dev/null +++ b/mitmproxy-rs/src/contentview.rs @@ -0,0 +1,31 @@ +use anyhow::Result; +use mitmproxy::contentviews::Contentview; +use pyo3::prelude::*; + +#[pyclass] +pub struct PyContentview(&'static dyn Contentview); + +impl PyContentview { + pub fn new<'py>( + py: Python<'py>, + contentview: &'static dyn Contentview, + ) -> PyResult> { + PyContentview(contentview).into_pyobject(py) + } +} + +#[pymethods] +impl PyContentview { + #[getter] + pub fn name(&self) -> &str { + self.0.name() + } + + pub fn deserialize<'py>(&self, data: Vec) -> Result { + self.0.deserialize(data) + } + + fn __repr__(&self) -> PyResult { + Ok(format!("<{} Contentview>", self.0.name())) + } +} diff --git a/mitmproxy-rs/src/lib.rs b/mitmproxy-rs/src/lib.rs index bde795bd..0b85a9c3 100644 --- a/mitmproxy-rs/src/lib.rs +++ b/mitmproxy-rs/src/lib.rs @@ -5,6 +5,7 @@ use std::sync::RwLock; use once_cell::sync::Lazy; use pyo3::{exceptions::PyException, prelude::*}; +mod contentview; mod dns_resolver; mod process_info; mod server; @@ -81,6 +82,18 @@ mod mitmproxy_rs { use crate::util::{genkey, pubkey}; } + #[pymodule] + mod contentviews { + use super::*; + use crate::contentview::PyContentview; + use mitmproxy::contentviews::*; + + #[pymodule_init] + fn init(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add("hex", PyContentview::new(m.py(), &HexStream())?) + } + } + #[pymodule_export] use crate::stream::Stream; diff --git a/src/contentviews/mod.rs b/src/contentviews/mod.rs new file mode 100644 index 00000000..c8055091 --- /dev/null +++ b/src/contentviews/mod.rs @@ -0,0 +1,81 @@ +use anyhow::Result; +use pretty_hex::{HexConfig, PrettyHex}; +use std::num::ParseIntError; + +#[derive(Debug)] +pub enum SerializeError { + InvalidFormat(String), +} + +pub trait Contentview: Send + Sync { + fn name(&self) -> &str; + fn deserialize(&self, data: Vec) -> Result; +} + +pub trait SerializableContentview: Contentview { + fn serialize(&self, data: String) -> Result, SerializeError>; +} + +#[derive(Default)] +pub struct HexStream(); + +impl Contentview for HexStream { + fn name(&self) -> &str { + "HexStream" + } + + fn deserialize(&self, data: Vec) -> Result { + Ok(data + .hex_conf(HexConfig { + title: false, + ascii: false, + width: 0, + group: 0, + chunk: 0, + max_bytes: usize::MAX, + display_offset: 0, + }) + .to_string()) + } +} + +impl SerializableContentview for HexStream { + fn serialize(&self, data: String) -> Result, SerializeError> { + (0..data.len()) + .step_by(2) + .map(|i| u8::from_str_radix(&data[i..i + 2], 16)) + .collect::, ParseIntError>>() + .map_err(|e| { + SerializeError::InvalidFormat(format!("Failed to parse hex string: {}", e)) + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_hexstream_deserialize() { + let hex_stream = HexStream::default(); + let data = b"foo".to_vec(); + let result = hex_stream.deserialize(data).unwrap(); + assert_eq!(result, "666f6f"); + } + + #[test] + fn test_hexstream_deserialize_empty() { + let hex_stream = HexStream::default(); + let data = vec![]; + let result = hex_stream.deserialize(data).unwrap(); + assert_eq!(result, ""); + } + + #[test] + fn test_hexstream_serialize() { + let hex_stream = HexStream::default(); + let data = "666f6f".to_string(); + let result = hex_stream.serialize(data).unwrap(); + assert_eq!(result, b"foo"); + } +} diff --git a/src/lib.rs b/src/lib.rs index ae3a8748..437e1cf7 100755 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,6 +1,7 @@ pub use network::MAX_PACKET_SIZE; pub mod certificates; +pub mod contentviews; pub mod dns; pub mod intercept_conf; pub mod ipc;