good_mitm_rule/
lib.rs

1pub use action::Action;
2pub use filter::Filter;
3pub use handler::*;
4use hyper::{header, header::HeaderValue, Body, Request, Response, StatusCode};
5use log::*;
6use mitm_core::mitm::RequestOrResponse;
7use std::vec::Vec;
8
9mod action;
10mod cache;
11mod filter;
12mod handler;
13
14#[derive(Debug, Clone)]
15pub struct Rule {
16    pub filters: Vec<Filter>,
17    pub actions: Vec<Action>,
18
19    pub url: Option<String>,
20}
21
22impl Rule {
23    pub async fn do_req(&mut self, req: Request<Body>) -> RequestOrResponse {
24        let url = req.uri().to_string();
25        self.url = Some(url.clone());
26        let mut tmp_req = req;
27
28        for action in &self.actions {
29            match action {
30                Action::Reject => {
31                    info!("[Reject] {}", url);
32                    let res = Response::builder()
33                        .status(StatusCode::BAD_GATEWAY)
34                        .body(Body::default())
35                        .unwrap();
36
37                    return RequestOrResponse::Response(res);
38                }
39
40                Action::Redirect(target) => {
41                    if target.contains('$') {
42                        for filter in self.filters.clone() {
43                            if let Filter::UrlRegex(re) = filter {
44                                let target = cache::get_regex(&re)
45                                    .replace(tmp_req.uri().to_string().as_str(), target.as_str())
46                                    .to_string();
47                                if let Ok(target_url) = HeaderValue::from_str(target.as_str()) {
48                                    let mut res = Response::builder()
49                                        .status(StatusCode::FOUND)
50                                        .body(Body::default())
51                                        .unwrap();
52                                    res.headers_mut().insert(header::LOCATION, target_url);
53                                    info!("[Redirect] {} -> {}", url, target);
54                                    return RequestOrResponse::Response(res);
55                                }
56                            }
57                        }
58                    }
59                    if let Ok(target_url) = HeaderValue::from_str(target.as_str()) {
60                        let mut res = Response::builder()
61                            .status(StatusCode::FOUND)
62                            .body(Body::default())
63                            .unwrap();
64                        res.headers_mut().insert(header::LOCATION, target_url);
65                        info!("[Redirect] {} -> {}", url, target);
66                        return RequestOrResponse::Response(res);
67                    };
68                }
69
70                Action::ModifyRequest(modify) => {
71                    info!("[ModifyRequest] {}", url);
72                    match modify.modify_req(tmp_req).await {
73                        Some(new_req) => tmp_req = new_req,
74                        None => {
75                            return RequestOrResponse::Response(
76                                Response::builder()
77                                    .status(StatusCode::BAD_REQUEST)
78                                    .body(Body::default())
79                                    .unwrap(),
80                            );
81                        }
82                    }
83                }
84
85                Action::LogReq => {
86                    info!("[LogRequest] {}", url);
87                    action::log_req(&tmp_req).await;
88                }
89
90                #[cfg(feature = "js")]
91                Action::Js(ref code) => {
92                    info!("[LogRequest] {}", url);
93                    if let Ok(req) = action::js::modify_req(code, tmp_req).await {
94                        return RequestOrResponse::Request(req);
95                    } else {
96                        return RequestOrResponse::Response(
97                            Response::builder()
98                                .status(StatusCode::BAD_REQUEST)
99                                .body(Body::default())
100                                .unwrap(),
101                        );
102                    }
103                }
104                _ => {}
105            }
106        }
107
108        RequestOrResponse::Request(tmp_req)
109    }
110
111    pub async fn do_res(&self, res: Response<Body>) -> Response<Body> {
112        let url = self.url.clone().unwrap_or_default();
113        let mut tmp_res = res;
114
115        for action in &self.actions {
116            match action {
117                Action::ModifyResponse(modify) => {
118                    info!("[ModifyResponse] {}", url);
119                    tmp_res = modify.modify_res(tmp_res).await
120                }
121                Action::LogRes => {
122                    info!("[LogResponse] {}", url);
123                    action::log_res(&tmp_res).await;
124                }
125
126                #[cfg(feature = "js")]
127                Action::Js(ref code) => {
128                    info!("[LogResponse] {}", url);
129                    if let Ok(res) = action::js::modify_res(code, tmp_res).await {
130                        return res;
131                    } else {
132                        return Response::builder()
133                            .status(StatusCode::BAD_REQUEST)
134                            .body(Body::default())
135                            .unwrap();
136                    }
137                }
138                _ => {}
139            }
140        }
141
142        tmp_res
143    }
144}
good_mitm_rule/lib.rs

good_mitm_rule/
lib.rs
