make contentviews interactive

mhils · mhils · commit 4a74daca7855 · 2025-03-14T21:16:54.000+01:00
diff --git a/mitmproxy-rs/mitmproxy_rs/__init__.pyi b/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",
diff --git a/mitmproxy-rs/mitmproxy_rs/contentviews.pyi b/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",
+]
diff --git a/mitmproxy-rs/src/contentview.rs b/mitmproxy-rs/src/contentview.rs
@@ -1,31 +1,65 @@
-use anyhow::Result;
-use mitmproxy::contentviews::Contentview;
+use anyhow::{anyhow, Result};
+use mitmproxy::contentviews::{Prettify, Reencode};
 use pyo3::prelude::*;
 
-#[pyclass]
-pub struct PyContentview(&'static dyn Contentview);
+#[pyclass(frozen, module = "mitmproxy_rs.contentviews", subclass)]
+pub struct Contentview(&'static dyn Prettify);
 
-impl PyContentview {
+impl Contentview {
     pub fn new<'py>(
         py: Python<'py>,
-        contentview: &'static dyn Contentview,
+        contentview: &'static dyn Prettify,
     ) -> PyResult<Bound<'py, Self>> {
-        PyContentview(contentview).into_pyobject(py)
+        Contentview(contentview).into_pyobject(py)
     }
 }
 
 #[pymethods]
-impl PyContentview {
+impl Contentview {
+    /// The name of this contentview.
     #[getter]
     pub fn name(&self) -> &str {
         self.0.name()
     }
 
+    /// Pretty-print an (encoded) message.
     pub fn deserialize<'py>(&self, data: Vec<u8>) -> Result<String> {
-        self.0.deserialize(data)
+        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()
+        ))
     }
 }
diff --git a/mitmproxy-rs/src/lib.rs b/mitmproxy-rs/src/lib.rs
@@ -2,6 +2,8 @@ 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::*};
 
@@ -85,12 +87,17 @@ mod mitmproxy_rs {
     #[pymodule]
     mod contentviews {
         use super::*;
-        use crate::contentview::PyContentview;
-        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(())
         }
     }
 
@@ -118,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)
+    }
+}
diff --git a/src/contentviews/hex_dump.rs b/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, "");
+    }
+}
diff --git a/src/contentviews/hex_stream.rs b/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");
+    }
+}
diff --git a/src/contentviews/mod.rs b/src/contentviews/mod.rs
@@ -1,81 +1,52 @@
+mod hex_dump;
+mod hex_stream;
+
 use anyhow::Result;
-use pretty_hex::{HexConfig, PrettyHex};
-use std::num::ParseIntError;
+use std::fmt::{Display, Formatter};
+
+pub use hex_dump::HexDump;
+pub use hex_stream::HexStream;
 
 #[derive(Debug)]
-pub enum SerializeError {
+pub enum ReencodeError {
     InvalidFormat(String),
 }
 
-pub trait Contentview: Send + Sync {
-    fn name(&self) -> &str;
-    fn deserialize(&self, data: Vec<u8>) -> Result<String>;
-}
-
-pub trait SerializableContentview: Contentview {
-    fn serialize(&self, data: String) -> Result<Vec<u8>, SerializeError>;
-}
-
-#[derive(Default)]
-pub struct HexStream();
-
-impl Contentview for HexStream {
-    fn name(&self) -> &str {
-        "HexStream"
+impl Display for ReencodeError {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match self {
+            ReencodeError::InvalidFormat(e) => {
+                write!(f, "invalid format: {}", e)
+            }
+        }
     }
+}
 
-    fn deserialize(&self, data: Vec<u8>) -> Result<String> {
-        Ok(data
-            .hex_conf(HexConfig {
-                title: false,
-                ascii: false,
-                width: 0,
-                group: 0,
-                chunk: 0,
-                max_bytes: usize::MAX,
-                display_offset: 0,
-            })
-            .to_string())
-    }
+#[derive(Debug)]
+pub enum PrettifyError {
+    Generic(String),
 }
 
-impl SerializableContentview for HexStream {
-    fn serialize(&self, data: String) -> Result<Vec<u8>, SerializeError> {
-        (0..data.len())
-            .step_by(2)
-            .map(|i| u8::from_str_radix(&data[i..i + 2], 16))
-            .collect::<Result<Vec<u8>, ParseIntError>>()
-            .map_err(|e| {
-                SerializeError::InvalidFormat(format!("Failed to parse hex string: {}", e))
-            })
+impl Display for PrettifyError {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match self {
+            PrettifyError::Generic(e) => {
+                write!(f, "deserialize error: {}", e)
+            }
+        }
     }
 }
 
-#[cfg(test)]
-mod tests {
-    use super::*;
+pub trait Prettify: Send + Sync {
+    fn name(&self) -> &str;
 
-    #[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");
+    fn instance_name(&self) -> String {
+        self.name().to_lowercase().replace(" ", "_")
     }
 
-    #[test]
-    fn test_hexstream_deserialize_empty() {
-        let hex_stream = HexStream::default();
-        let data = vec![];
-        let result = hex_stream.deserialize(data).unwrap();
-        assert_eq!(result, "");
-    }
+    fn deserialize(&self, data: Vec<u8>) -> Result<String, PrettifyError>;
+}
 
-    #[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");
-    }
+pub trait Reencode: Send + Sync {
+    fn serialize(&self, data: String) -> Result<Vec<u8>, ReencodeError>;
 }
