1use crate::error::Result;
4use bytes::Bytes;
5use slinger::{Body, Request, Response};
6use std::fmt;
7use std::net::SocketAddr;
8use std::sync::Arc;
9use std::time::{SystemTime, UNIX_EPOCH};
10
11#[derive(Clone)]
14pub struct MitmRequest {
15 pub source: Option<SocketAddr>,
17 pub destination: String,
19 pub timestamp: u64,
21 is_http: bool,
23 pub request: Request,
25}
26
27impl MitmRequest {
28 pub fn new(destination: impl Into<String>, request: Request) -> Self {
30 Self {
31 source: None,
32 destination: destination.into(),
33 timestamp: SystemTime::now()
34 .duration_since(UNIX_EPOCH)
35 .map(|d| d.as_millis() as u64)
36 .unwrap_or(0),
37 is_http: true,
38 request,
39 }
40 }
41
42 pub fn with_source(source: SocketAddr, destination: impl Into<String>, request: Request) -> Self {
44 Self {
45 source: Some(source),
46 destination: destination.into(),
47 timestamp: SystemTime::now()
48 .duration_since(UNIX_EPOCH)
49 .map(|d| d.as_millis() as u64)
50 .unwrap_or(0),
51 is_http: true,
52 request,
53 }
54 }
55
56 pub fn raw_tcp(destination: impl Into<String>, body: impl Into<Bytes>) -> Self {
58 let request = Request {
59 body: Some(Body::from(body.into())),
60 ..Default::default()
61 };
62 Self {
63 source: None,
64 destination: destination.into(),
65 timestamp: SystemTime::now()
66 .duration_since(UNIX_EPOCH)
67 .map(|d| d.as_millis() as u64)
68 .unwrap_or(0),
69 is_http: false,
70 request,
71 }
72 }
73
74 pub fn raw_tcp_with_source(
76 source: SocketAddr,
77 destination: impl Into<String>,
78 body: impl Into<Bytes>,
79 ) -> Self {
80 let request = Request {
81 body: Some(Body::from(body.into())),
82 ..Default::default()
83 };
84 Self {
85 source: Some(source),
86 destination: destination.into(),
87 timestamp: SystemTime::now()
88 .duration_since(UNIX_EPOCH)
89 .map(|d| d.as_millis() as u64)
90 .unwrap_or(0),
91 is_http: false,
92 request,
93 }
94 }
95
96 pub fn source(&self) -> Option<SocketAddr> {
98 self.source
99 }
100
101 pub fn destination(&self) -> &str {
103 &self.destination
104 }
105
106 pub fn timestamp(&self) -> u64 {
108 self.timestamp
109 }
110
111 pub fn request(&self) -> &Request {
113 &self.request
114 }
115
116 pub fn request_mut(&mut self) -> &mut Request {
118 &mut self.request
119 }
120
121 pub fn body(&self) -> Option<&Body> {
123 self.request.body.as_ref()
124 }
125
126 pub fn set_body(&mut self, body: impl Into<Bytes>) {
128 self.request.body = Some(Body::from(body.into()));
129 }
130
131 pub fn is_http(&self) -> bool {
133 self.is_http
134 }
135}
136
137impl fmt::Debug for MitmRequest {
138 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
139 f.debug_struct("MitmRequest")
140 .field("source", &self.source)
141 .field("destination", &self.destination)
142 .field("timestamp", &self.timestamp)
143 .field("is_http", &self.is_http())
144 .field("request", &self.request)
145 .finish()
146 }
147}
148
149#[derive(Clone)]
152pub struct MitmResponse {
153 pub source: String,
155 pub destination: Option<SocketAddr>,
157 pub timestamp: u64,
159 is_http: bool,
161 pub response: Response,
163}
164
165impl MitmResponse {
166 pub fn new(source: impl Into<String>, response: Response) -> Self {
168 Self {
169 source: source.into(),
170 destination: None,
171 timestamp: SystemTime::now()
172 .duration_since(UNIX_EPOCH)
173 .map(|d| d.as_millis() as u64)
174 .unwrap_or(0),
175 is_http: true,
176 response,
177 }
178 }
179
180 pub fn with_destination(
182 source: impl Into<String>,
183 destination: SocketAddr,
184 response: Response,
185 ) -> Self {
186 Self {
187 source: source.into(),
188 destination: Some(destination),
189 timestamp: SystemTime::now()
190 .duration_since(UNIX_EPOCH)
191 .map(|d| d.as_millis() as u64)
192 .unwrap_or(0),
193 is_http: true,
194 response,
195 }
196 }
197
198 pub fn raw_tcp(source: impl Into<String>, body: impl Into<Bytes>) -> Self {
200 let response = Response {
201 body: Some(Body::from(body.into())),
202 ..Default::default()
203 };
204 Self {
205 source: source.into(),
206 destination: None,
207 timestamp: SystemTime::now()
208 .duration_since(UNIX_EPOCH)
209 .map(|d| d.as_millis() as u64)
210 .unwrap_or(0),
211 is_http: false,
212 response,
213 }
214 }
215
216 pub fn raw_tcp_with_destination(
218 source: impl Into<String>,
219 destination: SocketAddr,
220 body: impl Into<Bytes>,
221 ) -> Self {
222 let response = Response {
223 body: Some(Body::from(body.into())),
224 ..Default::default()
225 };
226 Self {
227 source: source.into(),
228 destination: Some(destination),
229 timestamp: SystemTime::now()
230 .duration_since(UNIX_EPOCH)
231 .map(|d| d.as_millis() as u64)
232 .unwrap_or(0),
233 is_http: false,
234 response,
235 }
236 }
237
238 pub fn source(&self) -> &str {
240 &self.source
241 }
242
243 pub fn destination(&self) -> Option<SocketAddr> {
245 self.destination
246 }
247
248 pub fn timestamp(&self) -> u64 {
250 self.timestamp
251 }
252
253 pub fn response(&self) -> &Response {
255 &self.response
256 }
257
258 pub fn response_mut(&mut self) -> &mut Response {
260 &mut self.response
261 }
262
263 pub fn body(&self) -> Option<&Body> {
265 self.response.body.as_ref()
266 }
267
268 pub fn set_body(&mut self, body: impl Into<Bytes>) {
270 self.response.body = Some(Body::from(body.into()));
271 }
272
273 pub fn is_http(&self) -> bool {
275 self.is_http
276 }
277}
278
279impl fmt::Debug for MitmResponse {
280 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
281 f.debug_struct("MitmResponse")
282 .field("source", &self.source)
283 .field("destination", &self.destination)
284 .field("timestamp", &self.timestamp)
285 .field("is_http", &self.is_http())
286 .field("response", &self.response)
287 .finish()
288 }
289}
290
291#[async_trait::async_trait]
293pub trait RequestInterceptor: Send + Sync {
294 async fn intercept_request(&self, request: MitmRequest) -> Result<Option<MitmRequest>>;
298}
299
300#[async_trait::async_trait]
302pub trait ResponseInterceptor: Send + Sync {
303 async fn intercept_response(&self, response: MitmResponse) -> Result<Option<MitmResponse>>;
307}
308
309pub struct InterceptorHandler {
311 request_interceptors: Vec<Arc<dyn RequestInterceptor>>,
312 response_interceptors: Vec<Arc<dyn ResponseInterceptor>>,
313}
314
315impl InterceptorHandler {
316 pub fn new() -> Self {
318 Self {
319 request_interceptors: Vec::new(),
320 response_interceptors: Vec::new(),
321 }
322 }
323
324 pub fn add_request_interceptor(&mut self, interceptor: Arc<dyn RequestInterceptor>) {
326 self.request_interceptors.push(interceptor);
327 }
328
329 pub fn add_response_interceptor(&mut self, interceptor: Arc<dyn ResponseInterceptor>) {
331 self.response_interceptors.push(interceptor);
332 }
333
334 pub async fn process_request(&self, mut request: MitmRequest) -> Result<Option<MitmRequest>> {
336 for interceptor in &self.request_interceptors {
337 match interceptor.intercept_request(request).await? {
338 Some(modified) => request = modified,
339 None => return Ok(None), }
341 }
342 Ok(Some(request))
343 }
344
345 pub async fn process_response(&self, mut response: MitmResponse) -> Result<Option<MitmResponse>> {
347 for interceptor in &self.response_interceptors {
348 match interceptor.intercept_response(response).await? {
349 Some(modified) => response = modified,
350 None => return Ok(None), }
352 }
353 Ok(Some(response))
354 }
355
356 pub fn has_interceptors(&self) -> bool {
358 !self.request_interceptors.is_empty() || !self.response_interceptors.is_empty()
359 }
360}
361
362impl Default for InterceptorHandler {
363 fn default() -> Self {
364 Self::new()
365 }
366}
367
368pub struct Interceptor;
370
371impl Interceptor {
372 pub fn logging() -> LoggingInterceptor {
374 LoggingInterceptor
375 }
376}
377
378pub struct LoggingInterceptor;
380
381#[async_trait::async_trait]
382impl RequestInterceptor for LoggingInterceptor {
383 async fn intercept_request(&self, request: MitmRequest) -> Result<Option<MitmRequest>> {
384 if request.is_http() {
385 tracing::info!(
386 "[MITM] HTTP Request: {} {}",
387 request.request().method(),
388 request.request().uri()
389 );
390 for (name, value) in request.request().headers() {
391 tracing::info!(" {}: {:?}", name, value);
392 }
393 } else {
394 tracing::info!(
395 "[MITM] TCP Request to {}: {} bytes",
396 request.destination(),
397 request.body().map(|b| b.len()).unwrap_or(0)
398 );
399 }
400 if let Some(source) = request.source() {
401 tracing::info!(" From: {}", source);
402 }
403 tracing::info!(" Timestamp: {}", request.timestamp());
404 Ok(Some(request))
405 }
406}
407
408#[async_trait::async_trait]
409impl ResponseInterceptor for LoggingInterceptor {
410 async fn intercept_response(&self, response: MitmResponse) -> Result<Option<MitmResponse>> {
411 if response.is_http() {
412 tracing::info!(
413 "[MITM] HTTP Response: {}",
414 response.response().status_code()
415 );
416 for (name, value) in response.response().headers() {
417 tracing::info!(" {}: {:?}", name, value);
418 }
419 } else {
420 tracing::info!(
421 "[MITM] TCP Response from {}: {} bytes",
422 response.source(),
423 response.body().map(|b| b.len()).unwrap_or(0)
424 );
425 }
426 if let Some(destination) = response.destination() {
427 tracing::info!(" To: {}", destination);
428 }
429 tracing::info!(" Timestamp: {}", response.timestamp());
430 Ok(Some(response))
431 }
432}