1use crate::form_data::{analyze_form_data, filter_form_data};
2use crate::serializable::{SerializableRequest, SerializableResponse};
3use regex::Regex;
4use serde_json::{Map, Value};
5use std::collections::HashMap;
6use std::fmt::Debug;
7
8pub trait Filter: Debug + Send + Sync {
9 fn filter_request(&self, request: &mut SerializableRequest);
10 fn filter_response(&self, response: &mut SerializableResponse);
11}
12
13#[derive(Debug)]
14pub struct FilterChain {
15 filters: Vec<Box<dyn Filter>>,
16}
17
18impl FilterChain {
19 pub fn new() -> Self {
20 Self {
21 filters: Vec::new(),
22 }
23 }
24
25 pub fn add_filter(mut self, filter: Box<dyn Filter>) -> Self {
26 self.filters.push(filter);
27 self
28 }
29
30 pub fn filter_request(&self, request: &mut SerializableRequest) {
31 for filter in &self.filters {
32 filter.filter_request(request);
33 }
34 }
35
36 pub fn filter_response(&self, response: &mut SerializableResponse) {
37 for filter in &self.filters {
38 filter.filter_response(response);
39 }
40 }
41}
42
43impl Default for FilterChain {
44 fn default() -> Self {
45 Self::new()
46 }
47}
48
49#[derive(Debug)]
50pub struct HeaderFilter {
51 headers_to_remove: Vec<String>,
52 headers_to_replace: HashMap<String, String>,
53}
54
55impl HeaderFilter {
56 pub fn new() -> Self {
57 Self {
58 headers_to_remove: Vec::new(),
59 headers_to_replace: HashMap::new(),
60 }
61 }
62
63 pub fn remove_header(mut self, header: impl Into<String>) -> Self {
64 self.headers_to_remove.push(header.into());
65 self
66 }
67
68 pub fn replace_header(
69 mut self,
70 header: impl Into<String>,
71 replacement: impl Into<String>,
72 ) -> Self {
73 self.headers_to_replace
74 .insert(header.into(), replacement.into());
75 self
76 }
77
78 pub fn remove_auth_headers(self) -> Self {
79 self.remove_header("Authorization")
80 .remove_header("Cookie")
81 .remove_header("Set-Cookie")
82 .remove_header("X-API-Key")
83 .remove_header("X-Auth-Token")
84 }
85
86 fn filter_headers(&self, headers: &mut HashMap<String, Vec<String>>) {
87 for header in &self.headers_to_remove {
88 headers.remove(header);
89 }
90
91 for (header, replacement) in &self.headers_to_replace {
92 if let Some(values) = headers.get_mut(header) {
93 values.clear();
94 values.push(replacement.clone());
95 }
96 }
97 }
98}
99
100impl Filter for HeaderFilter {
101 fn filter_request(&self, request: &mut SerializableRequest) {
102 self.filter_headers(&mut request.headers);
103 }
104
105 fn filter_response(&self, response: &mut SerializableResponse) {
106 self.filter_headers(&mut response.headers);
107 }
108}
109
110impl Default for HeaderFilter {
111 fn default() -> Self {
112 Self::new()
113 }
114}
115
116#[derive(Debug)]
117pub struct BodyFilter {
118 json_keys_to_remove: Vec<String>,
119 json_keys_to_replace: HashMap<String, String>,
120 regex_replacements: Vec<(Regex, String)>,
121}
122
123impl BodyFilter {
124 pub fn new() -> Self {
125 Self {
126 json_keys_to_remove: Vec::new(),
127 json_keys_to_replace: HashMap::new(),
128 regex_replacements: Vec::new(),
129 }
130 }
131
132 pub fn remove_json_key(mut self, key: impl Into<String>) -> Self {
133 self.json_keys_to_remove.push(key.into());
134 self
135 }
136
137 pub fn replace_json_key(
138 mut self,
139 key: impl Into<String>,
140 replacement: impl Into<String>,
141 ) -> Self {
142 self.json_keys_to_replace
143 .insert(key.into(), replacement.into());
144 self
145 }
146
147 pub fn replace_regex(
148 mut self,
149 pattern: &str,
150 replacement: impl Into<String>,
151 ) -> Result<Self, regex::Error> {
152 let regex = Regex::new(pattern)?;
153 self.regex_replacements.push((regex, replacement.into()));
154 Ok(self)
155 }
156
157 pub fn remove_common_sensitive_keys(self) -> Self {
158 self.remove_json_key("password")
159 .remove_json_key("token")
160 .remove_json_key("api_key")
161 .remove_json_key("secret")
162 .remove_json_key("access_token")
163 .remove_json_key("refresh_token")
164 .remove_json_key("client_secret")
165 }
166
167 fn filter_json_object(&self, obj: &mut Map<String, Value>) {
168 for key in &self.json_keys_to_remove {
169 obj.remove(key);
170 }
171
172 for (key, replacement) in &self.json_keys_to_replace {
173 if obj.contains_key(key) {
174 obj.insert(key.clone(), Value::String(replacement.clone()));
175 }
176 }
177
178 for (_, value) in obj.iter_mut() {
179 self.filter_json_value(value);
180 }
181 }
182
183 fn filter_json_value(&self, value: &mut Value) {
184 match value {
185 Value::Object(obj) => self.filter_json_object(obj),
186 Value::Array(arr) => {
187 for item in arr.iter_mut() {
188 self.filter_json_value(item);
189 }
190 }
191 _ => {}
192 }
193 }
194
195 fn filter_body(&self, body: &mut Option<String>) {
196 if let Some(body_str) = body {
197 if let Ok(mut json_value) = serde_json::from_str::<Value>(body_str) {
198 self.filter_json_value(&mut json_value);
200 if let Ok(filtered_json) = serde_json::to_string(&json_value) {
201 *body_str = filtered_json;
202 }
203 } else if body_str.contains('=') && body_str.contains('&') {
204 let filtered = filter_form_data(body_str, "[FILTERED]");
206 *body_str = filtered;
207 } else {
208 for (regex, replacement) in &self.regex_replacements {
210 *body_str = regex.replace_all(body_str, replacement).to_string();
211 }
212 }
213 }
214 }
215}
216
217impl Filter for BodyFilter {
218 fn filter_request(&self, request: &mut SerializableRequest) {
219 self.filter_body(&mut request.body);
220 }
221
222 fn filter_response(&self, response: &mut SerializableResponse) {
223 self.filter_body(&mut response.body);
224 }
225}
226
227impl Default for BodyFilter {
228 fn default() -> Self {
229 Self::new()
230 }
231}
232
233#[derive(Debug)]
234pub struct UrlFilter {
235 query_params_to_remove: Vec<String>,
236 query_params_to_replace: HashMap<String, String>,
237}
238
239impl UrlFilter {
240 pub fn new() -> Self {
241 Self {
242 query_params_to_remove: Vec::new(),
243 query_params_to_replace: HashMap::new(),
244 }
245 }
246
247 pub fn remove_query_param(mut self, param: impl Into<String>) -> Self {
248 self.query_params_to_remove.push(param.into());
249 self
250 }
251
252 pub fn replace_query_param(
253 mut self,
254 param: impl Into<String>,
255 replacement: impl Into<String>,
256 ) -> Self {
257 self.query_params_to_replace
258 .insert(param.into(), replacement.into());
259 self
260 }
261
262 pub fn remove_common_sensitive_params(self) -> Self {
263 self.remove_query_param("api_key")
264 .remove_query_param("token")
265 .remove_query_param("access_token")
266 .remove_query_param("key")
267 }
268}
269
270impl Filter for UrlFilter {
271 fn filter_request(&self, request: &mut SerializableRequest) {
272 if let Ok(mut url) = url::Url::parse(&request.url) {
273 let mut query_pairs: Vec<(String, String)> = url
274 .query_pairs()
275 .map(|(k, v)| (k.to_string(), v.to_string()))
276 .collect();
277
278 query_pairs.retain(|(key, _)| !self.query_params_to_remove.contains(key));
279
280 for (key, value) in &mut query_pairs {
281 if let Some(replacement) = self.query_params_to_replace.get(key) {
282 *value = replacement.clone();
283 }
284 }
285
286 url.query_pairs_mut().clear();
287 for (key, value) in query_pairs {
288 url.query_pairs_mut().append_pair(&key, &value);
289 }
290
291 request.url = url.to_string();
292 }
293 }
294
295 fn filter_response(&self, _response: &mut SerializableResponse) {
296 }
298}
299
300impl Default for UrlFilter {
301 fn default() -> Self {
302 Self::new()
303 }
304}
305
306#[derive(Debug)]
307pub struct SmartFormFilter {
308 replacement_pattern: String,
309 verbose: bool,
310}
311
312impl SmartFormFilter {
313 pub fn new() -> Self {
314 Self {
315 replacement_pattern: "[FILTERED]".to_string(),
316 verbose: false,
317 }
318 }
319
320 pub fn with_replacement_pattern(mut self, pattern: impl Into<String>) -> Self {
321 self.replacement_pattern = pattern.into();
322 self
323 }
324
325 pub fn verbose(mut self) -> Self {
326 self.verbose = true;
327 self
328 }
329
330 fn filter_form_body(&self, body: &mut Option<String>) {
331 if let Some(body_str) = body {
332 if body_str.contains('=') && (body_str.contains('&') || !body_str.contains(' ')) {
334 if self.verbose {
335 println!("🔍 Analyzing form data in request body...");
336 let analysis = analyze_form_data(body_str);
337 analysis.print_summary();
338 }
339
340 let filtered = filter_form_data(body_str, &self.replacement_pattern);
341 *body_str = filtered;
342
343 if self.verbose {
344 println!("✅ Form data filtered");
345 }
346 }
347 }
348 }
349}
350
351impl Filter for SmartFormFilter {
352 fn filter_request(&self, request: &mut SerializableRequest) {
353 self.filter_form_body(&mut request.body);
354 }
355
356 fn filter_response(&self, _response: &mut SerializableResponse) {
357 }
359}
360
361impl Default for SmartFormFilter {
362 fn default() -> Self {
363 Self::new()
364 }
365}
366
367#[derive(Debug)]
368pub struct CustomFilter<F>
369where
370 F: Fn(&mut SerializableRequest, &mut SerializableResponse) + Send + Sync + Debug,
371{
372 filter_fn: F,
373}
374
375impl<F> CustomFilter<F>
376where
377 F: Fn(&mut SerializableRequest, &mut SerializableResponse) + Send + Sync + Debug,
378{
379 pub fn new(filter_fn: F) -> Self {
380 Self { filter_fn }
381 }
382}
383
384impl<F> Filter for CustomFilter<F>
385where
386 F: Fn(&mut SerializableRequest, &mut SerializableResponse) + Send + Sync + Debug,
387{
388 fn filter_request(&self, request: &mut SerializableRequest) {
389 let mut dummy_response = SerializableResponse {
390 status: 200,
391 headers: HashMap::new(),
392 body: None,
393 body_base64: None,
394 version: "Http1_1".to_string(),
395 };
396 (self.filter_fn)(request, &mut dummy_response);
397 }
398
399 fn filter_response(&self, response: &mut SerializableResponse) {
400 let mut dummy_request = SerializableRequest {
401 method: "GET".to_string(),
402 url: "https://example.com".to_string(),
403 headers: HashMap::new(),
404 body: None,
405 body_base64: None,
406 version: "Http1_1".to_string(),
407 };
408 (self.filter_fn)(&mut dummy_request, response);
409 }
410}