feat: Add log action

zu1k · Oct 11, 2021 · 79378dd · 79378dd
1 parent 6b13ceb
commit 79378dd
Show file tree

Hide file tree

Showing 6 changed files with 97 additions and 6 deletions.
diff --git a/rules/log.yaml b/rules/log.yaml
@@ -0,0 +1,5 @@
+- name: "log"
+  filter: all
+  action:
+    - log-req
+    - log-res
diff --git a/src/rule/action/log.rs b/src/rule/action/log.rs
@@ -0,0 +1,61 @@
+use http_mitm::hyper::{Body, Request, Response};
+use log::info;
+use std::fmt::Write;
+
+pub async fn log_req(req: &Request<Body>) {
+    let headers = req.headers();
+    let mut header_formated = String::new();
+    for (key, value) in headers {
+        let v = match value.to_str() {
+            Ok(v) => v.to_string(),
+            Err(_) => {
+                format!("[u8]; {}", value.len())
+            }
+        };
+        write!(
+            &mut header_formated,
+            "\t{:<20}{}\r\n",
+            format!("{}:", key.as_str()),
+            v
+        )
+        .unwrap();
+    }
+
+    info!(
+        "{} {}
+Headers:
+{}",
+        req.method(),
+        req.uri().to_string(),
+        header_formated
+    )
+}
+
+pub async fn log_res(res: &Response<Body>) {
+    let headers = res.headers();
+    let mut header_formated = String::new();
+    for (key, value) in headers {
+        let v = match value.to_str() {
+            Ok(v) => v.to_string(),
+            Err(_) => {
+                format!("[u8]; {}", value.len())
+            }
+        };
+        write!(
+            &mut header_formated,
+            "\t{:<20}{}\r\n",
+            format!("{}:", key.as_str()),
+            v
+        )
+        .unwrap();
+    }
+
+    info!(
+        "{} {:?}
+Headers:
+{}",
+        res.status(),
+        res.version(),
+        header_formated
+    )
+}
diff --git a/src/rule/action/mod.rs b/src/rule/action/mod.rs
@@ -1,6 +1,9 @@
 mod modify;
 use modify::*;
 
+mod log;
+pub use self::log::*;
+
 use serde::{Deserialize, Serialize};
 
 #[derive(Debug, Clone, Deserialize, Serialize)]
@@ -10,4 +13,6 @@ pub enum Action {
     Redirect(String),
     ModifyRequest(Modify),
     ModifyResponse(Modify),
+    LogRes,
+    LogReq,
 }
diff --git a/src/rule/file.rs b/src/rule/file.rs
@@ -12,6 +12,7 @@ pub struct Rule {
 #[derive(Debug, Clone, Deserialize, Serialize)]
 #[serde(rename_all = "kebab-case")]
 pub enum Filter {
+    All,
     Domain(String),
     DomainKeyword(String),
     DomainPrefix(String),

diff --git a/src/rule/filter/mod.rs b/src/rule/filter/mod.rs
@@ -5,6 +5,7 @@ pub use mitm_filter::*;
 
 #[derive(Debug, Clone)]
 pub enum Filter {
+    All,
     Domain(String),
     DomainKeyword(String),
     DomainPrefix(String),
@@ -14,6 +15,10 @@ pub enum Filter {
 
 #[allow(dead_code)]
 impl Filter {
+    pub fn new_all() -> Self {
+        Self::All
+    }
+
     pub fn new_domain(s: &str) -> Self {
         Self::Domain(s.to_lowercase())
     }
@@ -38,6 +43,7 @@ impl Filter {
     pub fn is_match_req(&self, req: &Request<Body>) -> bool {
         let host = req.uri().host().unwrap_or_default().to_lowercase();
         match self {
+            Self::All => true,
             Self::Domain(target) => host == *target,
             Self::DomainKeyword(target) => host.contains(target),
             Self::DomainPrefix(target) => host.starts_with(target),

diff --git a/src/rule/mod.rs b/src/rule/mod.rs
@@ -27,6 +27,7 @@ pub struct Rule {
 impl From<file::Rule> for Rule {
     fn from(rule: file::Rule) -> Self {
         let filter = match rule.filter {
+            file::Filter::All => Filter::new_all(),
             file::Filter::Domain(s) => Filter::new_domain(s.as_str()),
             file::Filter::DomainKeyword(s) => Filter::new_domain_keyword(s.as_str()),
             file::Filter::DomainPrefix(s) => Filter::new_domain_prefix(s.as_str()),
@@ -102,6 +103,11 @@ impl Rule {
                     }
                 }
 
+                Action::LogReq => {
+                    info!("[LogRequest] {}", url);
+                    action::log_req(&tmp_req).await;
+                }
+
                 _ => {}
             }
         }
@@ -114,12 +120,19 @@ impl Rule {
         let mut tmp_res = res;
 
         for action in &self.action {
-            if let Action::ModifyResponse(modify) = action {
-                info!("[ModifyResponse] {}", url);
-                if tmp_res.headers().get(header::CONTENT_ENCODING).is_some() {
-                    tmp_res = decode_response(tmp_res).unwrap()
-                };
-                tmp_res = modify.modify_res(tmp_res).await
+            match action {
+                Action::ModifyResponse(modify) => {
+                    info!("[ModifyResponse] {}", url);
+                    if tmp_res.headers().get(header::CONTENT_ENCODING).is_some() {
+                        tmp_res = decode_response(tmp_res).unwrap()
+                    };
+                    tmp_res = modify.modify_res(tmp_res).await
+                }
+                Action::LogRes => {
+                    info!("[LogResponse] {}", url);
+                    action::log_res(&tmp_res).await;
+                }
+                _ => {}
             }
         }