Add WireGuard handshake key logging

Cargo.toml

@@ -71,7 +71,7 @@ socket2 = "0.5.8"
 
 [patch.crates-io]
 # tokio = { path = "../tokio/tokio" }
-boringtun = { git = 'https://github.com/cloudflare/boringtun', rev = 'e3252d9c4f4c8fc628995330f45369effd4660a1' }
+boringtun = { git = "https://github.com/ntqbit/boringtun", branch = "feature/key_logger"}
 
 [target.'cfg(windows)'.dependencies.windows]
 version = "0.58.0"

mitmproxy-rs/src/server/wireguard.rs

@@ -6,7 +6,7 @@ use mitmproxy::packet_sources::wireguard::WireGuardConf;
 
 use pyo3::prelude::*;
 
-use boringtun::x25519::PublicKey;
+use boringtun::{noise::keys_logger::KeyLogger, x25519::PublicKey};
 
 use crate::server::base::Server;
 
@@ -60,7 +60,9 @@ impl WireGuardServer {
 /// - `peer_public_keys`: List of public X25519 keys for WireGuard peers as base64-encoded strings.
 /// - `handle_tcp_stream`: An async function that will be called for each new TCP `Stream`.
 /// - `handle_udp_stream`: An async function that will be called for each new UDP `Stream`.
+/// - `key_logger`: An optional function that will be called for each key when handshake is completed.
 #[pyfunction]
+#[pyo3(signature = (host, port, private_key, peer_public_keys, handle_tcp_stream, handle_udp_stream, key_logger=None))]
 pub fn start_wireguard_server(
     py: Python<'_>,
     host: String,
@@ -69,20 +71,37 @@ pub fn start_wireguard_server(
     peer_public_keys: Vec<String>,
     handle_tcp_stream: PyObject,
     handle_udp_stream: PyObject,
+    key_logger: Option<PyObject>,
 ) -> PyResult<Bound<PyAny>> {
     let private_key = string_to_key(private_key)?;
     let peer_public_keys = peer_public_keys
         .into_iter()
         .map(string_to_key)
         .collect::<PyResult<Vec<PublicKey>>>()?;
+
+    let key_logger = key_logger
+        .map(|key_logger| Box::new(PythonKeyLogger(key_logger)) as Box<dyn KeyLogger + Sync>);
+
     let conf = WireGuardConf {
         host,
         port,
         private_key,
         peer_public_keys,
+        key_logger,
     };
     pyo3_async_runtimes::tokio::future_into_py(py, async move {
         let (server, local_addr) = Server::init(conf, handle_tcp_stream, handle_udp_stream).await?;
         Ok(WireGuardServer { server, local_addr })
     })
 }
+
+struct PythonKeyLogger(PyObject);
+
+impl KeyLogger for PythonKeyLogger {
+    fn log_key(&self, name: &str, keymaterial: &str) {
+        Python::with_gil(|py| {
+            // The error is intentionally ignored.
+            let _ = self.0.call1(py, (name, keymaterial));
+        });
+    }
+}

src/packet_sources/wireguard.rs

@@ -8,6 +8,8 @@ use crate::messages::{
 use crate::network::{add_network_layer, MAX_PACKET_SIZE};
 use crate::packet_sources::{PacketSourceConf, PacketSourceTask};
 use anyhow::{anyhow, Context, Result};
+use boringtun::noise::handshake::HandshakeKeysListener;
+use boringtun::noise::keys_logger::{KeyLogger, KeysLogger};
 use boringtun::noise::{
     errors::WireGuardError, handshake::parse_handshake_anon, Packet, Tunn, TunnResult,
 };
@@ -40,6 +42,15 @@ pub struct WireGuardConf {
     pub port: u16,
     pub private_key: StaticSecret,
     pub peer_public_keys: Vec<PublicKey>,
+    pub key_logger: Option<Box<dyn KeyLogger + Sync>>,
+}
+
+struct KeyLoggerWrapper(Box<dyn KeyLogger>);
+
+impl KeyLogger for KeyLoggerWrapper {
+    fn log_key(&self, name: &str, keymaterial: &str) {
+        self.0.log_key(name, keymaterial);
+    }
 }
 
 impl PacketSourceConf for WireGuardConf {
@@ -59,13 +70,17 @@ impl PacketSourceConf for WireGuardConf {
         let (network_task_handle, net_tx, net_rx) =
             add_network_layer(transport_events_tx, transport_commands_rx, shutdown);
 
+        let handshake_keys_listener = self
+            .key_logger
+            .map(|key_logger| Arc::new(KeysLogger::new(KeyLoggerWrapper(key_logger))));
+
         // initialize WireGuard server
         let mut peers_by_idx = HashMap::new();
         let mut peers_by_key = HashMap::new();
         for public_key in self.peer_public_keys {
             let index = peers_by_idx.len() as u32;
 
-            let tunnel = Tunn::new(
+            let mut tunnel = Tunn::new(
                 self.private_key.clone(),
                 public_key,
                 None,
@@ -75,6 +90,13 @@ impl PacketSourceConf for WireGuardConf {
             )
             .map_err(|error| anyhow!(error))?;
 
+            // Set key logger, if any.
+            if let Some(keys_listener) = &handshake_keys_listener {
+                tunnel.set_handshake_keys_listener(
+                    Arc::clone(keys_listener) as Arc<dyn HandshakeKeysListener>
+                );
+            }
+
             let peer = Arc::new(Mutex::new(WireGuardPeer {
                 tunnel,
                 endpoint: None,