make contentviews interactive

Author: mhils

mitmproxy-rs/mitmproxy_rs/__init__.pyi

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from typing import Any, Literal
 from typing import final, overload, TypeVar
-from . import certs, dns, local, process_info, tun, udp, wireguard
+from . import certs, contentviews, dns, local, process_info, tun, udp, wireguard
 
 T = TypeVar("T")
 
@@ -57,6 +57,7 @@ class Stream:
 
 __all__ = [
     "certs",
+    "contentviews",
     "dns",
     "local",
     "process_info",

mitmproxy-rs/mitmproxy_rs/contentviews.pyi

@@ -0,0 +1,22 @@
+from typing import ClassVar, final
+
+class Contentview:
+    name: ClassVar[str]
+
+    def deserialize(self, data: bytes) -> str:
+        pass
+
+@final
+class InteractiveContentview(Contentview):
+    def serialize(self, data: str) -> bytes:
+        pass
+
+hex_dump: Contentview
+hex_stream: InteractiveContentview
+
+__all__ = [
+    "Contentview",
+    "InteractiveContentview",
+    "hex_dump",
+    "hex_stream",
+]

mitmproxy-rs/src/contentview.rs

@@ -1,32 +1,65 @@
+use anyhow::{anyhow, Result};
+use mitmproxy::contentviews::{Prettify, Reencode};
 use pyo3::prelude::*;
-use anyhow::{Result};
-use mitmproxy::contentviews::Contentview;
 
-#[pyclass]
-pub struct PyContentview(&'static dyn Contentview);
+#[pyclass(frozen, module = "mitmproxy_rs.contentviews", subclass)]
+pub struct Contentview(&'static dyn Prettify);
 
-impl PyContentview {
-    pub fn new<'py>(py: Python<'py>, contentview: &'static dyn Contentview) -> PyResult<Bound<'py, Self>> {
-        PyContentview(contentview).into_pyobject(py)
+impl Contentview {
+    pub fn new<'py>(
+        py: Python<'py>,
+        contentview: &'static dyn Prettify,
+    ) -> PyResult<Bound<'py, Self>> {
+        Contentview(contentview).into_pyobject(py)
     }
 }
 
 #[pymethods]
-impl PyContentview {
-
+impl Contentview {
+    /// The name of this contentview.
     #[getter]
     pub fn name(&self) -> &str {
         self.0.name()
     }
 
-    pub fn deserialize<'py>(
-        &self,
-        data: Vec<u8>
-    ) -> Result<String> {
-        self.0.deserialize(data)
+    /// Pretty-print an (encoded) message.
+    pub fn deserialize<'py>(&self, data: Vec<u8>) -> Result<String> {
+        self.0.deserialize(data).map_err(|e| anyhow!("{e}"))
     }
 
     fn __repr__(&self) -> PyResult<String> {
-        Ok(format!("<{} Contentview>", self.0.name()))
+        Ok(format!(
+            "<mitmproxy_rs.contentview.Contentview: {}>",
+            self.0.name()
+        ))
+    }
+}
+
+#[pyclass(frozen, module = "mitmproxy_rs.contentviews", extends=Contentview)]
+pub struct InteractiveContentview(&'static dyn Reencode);
+
+impl InteractiveContentview {
+    /// Argument passed twice because of https://github.com/rust-lang/rust/issues/65991
+    pub fn new<'py, T: Prettify + Reencode>(
+        py: Python<'py>,
+        cv: &'static T,
+    ) -> PyResult<Bound<'py, Self>> {
+        let cls =
+            PyClassInitializer::from(Contentview(cv)).add_subclass(InteractiveContentview(cv));
+        Bound::new(py, cls)
+    }
+}
+
+#[pymethods]
+impl InteractiveContentview {
+    pub fn serialize<'py>(&self, data: String) -> Result<Vec<u8>> {
+        self.0.serialize(data).map_err(|e| anyhow!("{e}"))
     }
-}
+
+    fn __repr__(self_: PyRef<'_, Self>) -> PyResult<String> {
+        Ok(format!(
+            "<mitmproxy_rs.contentview.InteractiveContentview: {}>",
+            self_.as_super().name()
+        ))
+    }
+}

mitmproxy-rs/src/lib.rs

@@ -2,17 +2,19 @@ extern crate core;
 
 use std::sync::RwLock;
 
+use crate::contentview::{Contentview, InteractiveContentview};
+use mitmproxy::contentviews::{Prettify, Reencode};
 use once_cell::sync::Lazy;
 use pyo3::{exceptions::PyException, prelude::*};
 
+mod contentview;
 mod dns_resolver;
 mod process_info;
 mod server;
 mod stream;
 pub mod task;
 mod udp_client;
 mod util;
-mod contentview;
 
 static LOGGER_INITIALIZED: Lazy<RwLock<bool>> = Lazy::new(|| RwLock::new(false));
 
@@ -84,16 +86,18 @@ mod mitmproxy_rs {
 
     #[pymodule]
     mod contentviews {
-        use crate::contentview::PyContentview;
         use super::*;
-        use mitmproxy::contentviews::*;
+        #[pymodule_export]
+        use crate::contentview::Contentview;
+        #[pymodule_export]
+        use crate::contentview::InteractiveContentview;
+        use mitmproxy::contentviews::{HexDump, HexStream};
 
         #[pymodule_init]
         fn init(m: &Bound<'_, PyModule>) -> PyResult<()> {
-            m.add(
-                "hex",
-                PyContentview::new(m.py(), &HexStream())?
-            )
+            m.add_contentview(&HexDump)?;
+            m.add_interactive_contentview(&HexStream)?;
+            Ok(())
         }
     }
 
@@ -121,3 +125,19 @@ mod mitmproxy_rs {
         Ok(())
     }
 }
+
+trait AddContentview {
+    fn add_contentview<T: Prettify>(&self, cv: &'static T) -> PyResult<()>;
+    fn add_interactive_contentview<T: Prettify + Reencode>(&self, i: &'static T) -> PyResult<()>;
+}
+
+impl AddContentview for Bound<'_, PyModule> {
+    fn add_contentview<T: Prettify>(&self, cv: &'static T) -> PyResult<()> {
+        let view = Contentview::new(self.py(), cv)?;
+        self.add(cv.instance_name(), view)
+    }
+    fn add_interactive_contentview<T: Prettify + Reencode>(&self, cv: &'static T) -> PyResult<()> {
+        let view = InteractiveContentview::new(self.py(), cv)?;
+        self.add(cv.instance_name(), view)
+    }
+}

src/contentviews/hex_dump.rs

@@ -0,0 +1,47 @@
+use crate::contentviews::{Prettify, PrettifyError};
+use pretty_hex::{HexConfig, PrettyHex};
+
+pub struct HexDump;
+
+impl Prettify for HexDump {
+    fn name(&self) -> &'static str {
+        "Hex Dump"
+    }
+
+    fn deserialize(&self, data: Vec<u8>) -> Result<String, PrettifyError> {
+        Ok(format!(
+            "{:?}",
+            data.hex_conf(HexConfig {
+                title: false,
+                ascii: true,
+                width: 16,
+                group: 4,
+                chunk: 1,
+                max_bytes: usize::MAX,
+                display_offset: 0,
+            })
+        ))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_hexdump_deserialize() {
+        let data = b"abcd".to_vec();
+        let result = HexDump.deserialize(data).unwrap();
+        assert_eq!(
+            result,
+            "0000:   61 62 63 64                                          abcd"
+        );
+    }
+
+    #[test]
+    fn test_hexdump_deserialize_empty() {
+        let data = vec![];
+        let result = HexDump.deserialize(data).unwrap();
+        assert_eq!(result, "");
+    }
+}

src/contentviews/hex_stream.rs

@@ -0,0 +1,61 @@
+use crate::contentviews::{Prettify, PrettifyError, Reencode, ReencodeError};
+use pretty_hex::{HexConfig, PrettyHex};
+use std::num::ParseIntError;
+
+pub struct HexStream;
+
+impl Prettify for HexStream {
+    fn name(&self) -> &'static str {
+        "Hex Stream"
+    }
+
+    fn deserialize(&self, data: Vec<u8>) -> Result<String, PrettifyError> {
+        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 Reencode for HexStream {
+    fn serialize(&self, data: String) -> anyhow::Result<Vec<u8>, ReencodeError> {
+        (0..data.len())
+            .step_by(2)
+            .map(|i| u8::from_str_radix(&data[i..i + 2], 16))
+            .collect::<anyhow::Result<Vec<u8>, ParseIntError>>()
+            .map_err(|e| ReencodeError::InvalidFormat(e.to_string()))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_hexstream_deserialize() {
+        let data = b"foo".to_vec();
+        let result = HexStream.deserialize(data).unwrap();
+        assert_eq!(result, "666f6f");
+    }
+
+    #[test]
+    fn test_hexstream_deserialize_empty() {
+        let data = vec![];
+        let result = HexStream.deserialize(data).unwrap();
+        assert_eq!(result, "");
+    }
+
+    #[test]
+    fn test_hexstream_serialize() {
+        let data = "666f6f".to_string();
+        let result = HexStream.serialize(data).unwrap();
+        assert_eq!(result, b"foo");
+    }
+}