Merge pull request #25 from th7nder/t-17

emanuele-em · web-flow · commit b2accb685362 · 2023-02-24T14:20:17.000+01:00
Add listening interface and port selection
diff --git a/mitm_proxy/src/main.rs b/mitm_proxy/src/main.rs
@@ -1,19 +1,13 @@
+mod managed_proxy;
 mod mitm_proxy;
 mod requests;
 
-use std::{
-    sync::mpsc::{ sync_channel},
-    thread, net::SocketAddr,
-};
-
 use crate::mitm_proxy::MitmProxy;
 
 use eframe::{
     egui::{self, CentralPanel, Vec2},
-    run_native, App
+    run_native, App,
 };
-use proxyapi::proxy::Proxy;
-use tokio::runtime::Runtime;
 
 static X: f32 = 980.;
 static Y: f32 = 960.0;
@@ -25,25 +19,18 @@ static PADDING: f32 = 20.;
 
 impl App for MitmProxy {
     fn update(&mut self, ctx: &egui::Context, frame: &mut eframe::Frame) {
-
         ctx.request_repaint();
 
         self.manage_theme(ctx);
 
         self.render_top_panel(ctx, frame);
-        
-        CentralPanel::default().show(ctx, |ui|{
-                self.render_columns(ui);
+
+        CentralPanel::default().show(ctx, |ui| {
+            self.render_columns(ui);
         });
     }
 }
 
-async fn shutdown_signal() {
-    tokio::signal::ctrl_c()
-        .await
-        .expect("Failed to install CTRL+C signal handler");
-}
-
 fn load_icon(path: &str) -> eframe::IconData {
     let (icon_rgba, icon_width, icon_height) = {
         let image = image::open(path)
@@ -66,24 +53,9 @@ fn main() {
     native_options.initial_window_size = Some(Vec2::new(X, Y));
     native_options.icon_data = Some(load_icon("./assets/logo.png"));
 
-    // create the app with listener false
-    // update listener when it is true
-
-    let (tx, rx) = sync_channel(1);
-    let rt = Runtime::new().unwrap();
-    let addr = SocketAddr::new([127,0,0,1].into(), 8100);
-
-    thread::spawn(move || {
-        rt.block_on( async move {
-                if let Err(e) = Proxy::new(addr, Some(tx.clone())).start(shutdown_signal()).await{
-                    eprintln!("Error running proxy on {:?}: {e}", addr);
-                }
-        })
-    });
-
     run_native(
         "Man In The Middle Proxy",
         native_options,
-        Box::new(|cc| Box::new(MitmProxy::new(cc, rx))),
+        Box::new(|cc| Box::new(MitmProxy::new(cc))),
     )
 }
diff --git a/mitm_proxy/src/managed_proxy.rs b/mitm_proxy/src/managed_proxy.rs
@@ -0,0 +1,65 @@
+use std::{
+    net::SocketAddr,
+    sync::mpsc::Receiver,
+    thread::{self, JoinHandle},
+};
+
+use proxyapi::{Proxy, ProxyHandler};
+use tokio::{runtime::Runtime, sync::oneshot::Sender};
+
+use crate::requests::RequestInfo;
+
+pub struct ManagedProxy {
+    rx: Receiver<ProxyHandler>,
+    close: Option<Sender<()>>,
+    thread: Option<JoinHandle<()>>,
+}
+
+impl ManagedProxy {
+    pub fn new(addr: SocketAddr) -> ManagedProxy {
+        let (tx, rx) = std::sync::mpsc::sync_channel(1);
+        let (close_tx, close_rx) = tokio::sync::oneshot::channel();
+
+        let rt = Runtime::new().unwrap();
+
+        let thread = thread::spawn(move || {
+            rt.block_on(async move {
+                if let Err(e) = Proxy::new(addr, Some(tx.clone()))
+                    .start(async move {
+                        let _ = close_rx.await;
+                    })
+                    .await
+                {
+                    eprintln!("Error running proxy on {:?}: {e}", addr);
+                }
+            })
+        });
+
+        ManagedProxy {
+            rx,
+            close: Some(close_tx),
+            thread: Some(thread),
+        }
+    }
+
+    pub fn try_recv_request(&mut self) -> Option<RequestInfo> {
+        match self.rx.try_recv() {
+            Ok(l) => {
+                let (request, response) = l.to_parts();
+                Some(RequestInfo::new(request, response))
+            }
+            _ => None,
+        }
+    }
+}
+
+impl Drop for ManagedProxy {
+    fn drop(&mut self) {
+        if let Some(t) = self.thread.take() {
+            if let Some(close) = self.close.take() {
+                let _ = close.send(());
+            }
+            t.join().expect("Couldn't gracefully shutdown the proxy.")
+        }
+    }
+}
diff --git a/mitm_proxy/src/mitm_proxy.rs b/mitm_proxy/src/mitm_proxy.rs
@@ -1,23 +1,21 @@
-use std::{
-    fmt::{Display},
-    sync::mpsc::Receiver,
-};
+use std::{fmt::Display, net::SocketAddr};
 
 use crate::{
+    managed_proxy::ManagedProxy,
     requests::{InfoOptions, RequestInfo},
     PADDING,
 };
 
 use eframe::{
     egui::{
         self, ComboBox, FontData, FontDefinitions, FontFamily, Grid, Layout, RichText, ScrollArea,
-        Style, TextStyle::*, TopBottomPanel, Visuals,
+        Style, TextEdit, TextStyle::*, TopBottomPanel, Visuals,
     },
-    epaint::FontId,
+    epaint::{Color32, FontId},
     Frame,
 };
 use egui_extras::{Column, TableBuilder};
-use proxyapi::{hyper::Method, *};
+use proxyapi::hyper::Method;
 use serde::{Deserialize, Serialize};
 
 #[derive(Serialize, Deserialize)]
@@ -77,6 +75,7 @@ struct MitmProxyState {
     selected_request: Option<usize>,
     selected_request_method: MethodFilter,
     detail_option: InfoOptions,
+    listen_on: String,
 }
 
 impl MitmProxyState {
@@ -85,6 +84,7 @@ impl MitmProxyState {
             selected_request: None,
             selected_request_method: MethodFilter::All,
             detail_option: InfoOptions::Request,
+            listen_on: "127.0.0.1:8100".to_string(),
         }
     }
 }
@@ -93,11 +93,11 @@ pub struct MitmProxy {
     requests: Vec<RequestInfo>,
     config: MitmProxyConfig,
     state: MitmProxyState,
-    rx: Receiver<ProxyHandler>,
+    proxy: Option<ManagedProxy>,
 }
 
 impl MitmProxy {
-    pub fn new(cc: &eframe::CreationContext<'_>, rx: Receiver<ProxyHandler>) -> Self {
+    pub fn new(cc: &eframe::CreationContext<'_>) -> Self {
         Self::configure_fonts(cc);
         let config: MitmProxyConfig = confy::load("MitmProxy", None).unwrap_or_default();
         let state = MitmProxyState::new();
@@ -106,7 +106,7 @@ impl MitmProxy {
             requests: vec![],
             config,
             state,
-            rx,
+            proxy: None,
         }
     }
 
@@ -117,6 +117,23 @@ impl MitmProxy {
         }
     }
 
+    fn start_proxy(&mut self, addr: SocketAddr) {
+        assert!(self.proxy.is_none());
+
+        self.proxy = Some(ManagedProxy::new(addr));
+        self.requests = vec![];
+    }
+
+    fn stop_proxy(&mut self) {
+        assert!(self.proxy.is_some());
+
+        self.proxy.take();
+    }
+
+    fn is_running(&self) -> bool {
+        return self.proxy.is_some();
+    }
+
     fn configure_fonts(cc: &eframe::CreationContext<'_>) {
         let mut fonts = FontDefinitions::default();
 
@@ -255,19 +272,14 @@ impl MitmProxy {
             });
     }
 
-    pub fn update_requests(&mut self) -> Option<RequestInfo> {
-        match self.rx.try_recv() {
-            Ok(l) => {
-                let (request,response) = l.to_parts();
-                Some(RequestInfo::new(request,response))
-            },
-            _ => None,
-        }
-    }
-
     pub fn render_columns(&mut self, ui: &mut egui::Ui) {
-        if let Some(request) = self.update_requests() {
-            self.requests.push(request);
+        if !self.is_running() {
+            return;
+        }
+        if let Some(ref mut proxy) = self.proxy {
+            if let Some(request) = proxy.try_recv_request() {
+                self.requests.push(request);
+            }
         }
 
         if let Some(i) = self.state.selected_request {
@@ -292,6 +304,37 @@ impl MitmProxy {
     pub fn render_top_panel(&mut self, ctx: &egui::Context, _frame: &mut Frame) {
         TopBottomPanel::top("top_panel").show(ctx, |ui| {
             ui.add_space(PADDING);
+            egui::menu::bar(ui, |ui| -> egui::InnerResponse<_> {
+                ui.with_layout(Layout::left_to_right(eframe::emath::Align::Min), |ui| {
+                    if !self.is_running() {
+                        TextEdit::singleline(&mut self.state.listen_on).show(ui);
+
+                        match self.state.listen_on.parse::<SocketAddr>() {
+                            Ok(addr) => {
+                                let start_button = ui.button("▶").on_hover_text("Start");
+                                if start_button.clicked() {
+                                    self.start_proxy(addr);
+                                }
+                            }
+                            Err(_err) => {
+                                ui.label(
+                                    RichText::new("Provided invalid IP address")
+                                        .color(Color32::RED),
+                                );
+                            }
+                        };
+                    } else {
+                        TextEdit::singleline(&mut self.state.listen_on)
+                            .interactive(false)
+                            .show(ui);
+                        let stop_button = ui.button("■").on_hover_text("Stop");
+                        if stop_button.clicked() {
+                            self.stop_proxy();
+                        }
+                    }
+                })
+            });
+
             egui::menu::bar(ui, |ui| -> egui::InnerResponse<_> {
                 ui.with_layout(Layout::left_to_right(eframe::emath::Align::Min), |ui| {
                     let clean_btn = ui.button("🚫").on_hover_text("Clear");
@@ -326,14 +369,12 @@ impl MitmProxy {
                 });
 
                 ui.with_layout(Layout::right_to_left(eframe::emath::Align::Min), |ui| {
-                    
                     let theme_btn = ui
                         .button(match self.config.dark_mode {
                             true => "🔆",
                             false => "🌙",
                         })
                         .on_hover_text("Toggle theme");
-                    
 
                     if theme_btn.clicked() {
                         self.config.dark_mode = !self.config.dark_mode
