add syntax highlight module

Author: mhils

Cargo.lock

@@ -73,6 +73,10 @@ serde_yaml = "0.9"
 rmp-serde = "1.1"
 protobuf = "3.7.2"
 regex = "1.10.3"
+tree-sitter-highlight = "0.25.3"
+tree-sitter-yaml = "0.7.0"
+tree-sitter-xml = "0.7.0"
+tree-sitter = "0.25.3"
 
 [patch.crates-io]
 # tokio = { path = "../tokio/tokio" }

Cargo.toml

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

mitmproxy-rs/mitmproxy_rs/__init__.pyi

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+def syntax_highlight(s: str, language: str) -> list[tuple[str, str]]:
+    pass
+
+def all_tags(language: str) -> list[str]:
+    pass
+
+__all__ = ["syntax_highlight", "all_tags"] 

mitmproxy-rs/mitmproxy_rs/syntax_highlight.pyi

@@ -12,6 +12,7 @@ mod dns_resolver;
 mod process_info;
 mod server;
 mod stream;
+mod syntax_highlight;
 pub mod task;
 mod udp_client;
 mod util;
@@ -126,6 +127,14 @@ mod mitmproxy_rs {
 
         Ok(())
     }
+
+    #[pymodule]
+    mod syntax_highlight {
+        #[pymodule_export]
+        use crate::syntax_highlight::highlight;
+        #[pymodule_export]
+        use crate::syntax_highlight::all_tags;
+    }
 }
 
 trait AddContentview {

mitmproxy-rs/src/lib.rs

@@ -186,7 +186,7 @@ impl Stream {
                 }
                 _ => (),
             },
-            TunnelInfo::None {} => (),
+            TunnelInfo::None => (),
         }
         match default {
             Some(x) => Ok(x),

mitmproxy-rs/src/stream.rs

@@ -0,0 +1,29 @@
+#[allow(unused_imports)]
+use anyhow::{anyhow, Result};
+
+use pyo3::{exceptions::PyValueError, prelude::*};
+
+fn str_to_language(s: &str) -> PyResult<mitmproxy::syntax_highlight::Language> {
+    match s {
+        "xml" => Ok(mitmproxy::syntax_highlight::Language::Xml),
+        "yaml" => Ok(mitmproxy::syntax_highlight::Language::Yaml),
+        other => Err(PyErr::new::<PyValueError, _>(format!(
+            "Unsupported language: {other}"
+        ))),
+    }
+}
+
+/// Transform a text into tagged chunks for text.
+#[pyfunction]
+pub fn highlight(s: String, language: &str) -> PyResult<Vec<(&'static str, String)>> {
+    let language = str_to_language(language)?;
+    language.highlight(s.as_bytes())
+        .map_err(|e| PyValueError::new_err(e.to_string()))
+}
+
+/// Return the list of all possible tags for a given language.
+#[pyfunction]
+pub fn all_tags(language: &str) -> PyResult<&[&str]> {
+    let language = str_to_language(language)?;
+    Ok(language.all_tags())
+}

mitmproxy-rs/src/syntax_highlight.rs

@@ -1,6 +1,6 @@
+use crate::contentviews::hex_stream::is_binary;
 use crate::contentviews::{Metadata, Prettify};
 use pretty_hex::{HexConfig, PrettyHex};
-use crate::contentviews::hex_stream::is_binary;
 
 pub struct HexDump;
 

src/contentviews/hex_dump.rs

@@ -10,5 +10,6 @@ pub mod network;
 pub mod packet_sources;
 pub mod processes;
 pub mod shutdown;
+pub mod syntax_highlight;
 #[cfg(windows)]
 pub mod windows;

src/lib.rs

@@ -0,0 +1,65 @@
+use super::Chunk;
+use anyhow::{Context, Result};
+use tree_sitter_highlight::{HighlightConfiguration, HighlightEvent, Highlighter};
+
+pub fn highlight(
+    language: tree_sitter::Language,
+    highlights_query: &str,
+    tags: &[&'static str],
+    input: &[u8],
+) -> Result<Vec<Chunk>> {
+    let mut highlighter = Highlighter::new();
+    let mut config = HighlightConfiguration::new(language, "", highlights_query, "", "")
+        .context("failed to create highlight configuration")?;
+    config.configure(tags);
+
+    let highlights = highlighter
+        .highlight(&config, input, None, |_| None)
+        .context("failed to highlight")?;
+
+    let mut chunks: Vec<Chunk> = Vec::new();
+    let mut tag: Option<&'static str> = None;
+
+    for event in highlights {
+        let event = event.context("highlighter failure")?;
+        match event {
+            HighlightEvent::Source { start, end } => {
+                let contents = &input[start..end];
+                let tag_str = tag.unwrap_or("");
+                
+                match chunks.last_mut() {
+                    Some(x) if x.0 == tag_str => {
+                        x.1.push_str(&String::from_utf8_lossy(contents));
+                    }
+                    _ => chunks.push(
+                        (tag_str, String::from_utf8_lossy(contents).to_string())
+                    ),
+                }
+            }
+            HighlightEvent::HighlightStart(s) => {
+                tag = Some(tags[s.0]);
+            }
+            HighlightEvent::HighlightEnd => {
+                tag = None;
+            }
+        }
+    }
+    Ok(chunks)
+}
+
+#[cfg(test)]
+pub(super) fn test_tags_ok(
+    language: tree_sitter::Language,
+    highlights_query: &str,
+    tags: &[&'static str],
+) {
+    let config = HighlightConfiguration::new(language, "", highlights_query, "", "").unwrap();
+    for &tag in tags {
+        assert!(
+            config.names().iter().any(|name| name.contains(tag)),
+            "Invalid tag: {},\nAllowed tags: {:?}",
+            tag,
+            config.names()
+        );
+    }
+}