Merge pull request #4 from AikidoSec/Add-route-builder

Add route_builder module with make_route_from_url

Cargo.toml

@@ -10,3 +10,7 @@ crate-type = ["dylib"]
 
 [dependencies]
 regex = "1.10.6"
+url = "2.5.2"
+
+[dev-dependencies]
+rand = "0.8.5"

README.md

@@ -5,13 +5,22 @@ Zen Rust library is a library that can be used via FFI in different languages. C
 ## Python FFI Example code : 
 ```py
 import ctypes
-
-zen_rustlib = ctypes.CDLL("target/release/libzen_rustlib.so.so")
-
+zen_rustlib = ctypes.CDLL("target/release/libzen_rustlib.so")
 
 if __name__ == "__main__":
     command = "whoami | shell".encode("utf-8")
     userinput = "whoami".encode("utf-8")
     result = zen_rustlib.detect_shell_injection(command, userinput)
     print("Result", bool(result))
 ```
+For build_route_from_url : 
+```py
+import ctypes
+zen_rustlib = ctypes.CDLL("target/release/libzen_rustlib.so")
+zen_rustlib.build_route_from_url.restype = ctypes.c_char_p
+
+if __name__ == "__main__":
+    url = "https://aikido.dev/1234/".encode("utf-8")
+    result = zen_rustlib.build_route_from_url(url)
+    print("Result", result.decode("utf-8"))
+```

src/helpers/mod.rs

@@ -0,0 +1,2 @@
+pub mod try_parse_url_path;
+pub mod try_parse_url_path_test;

src/helpers/try_parse_url_path.rs

@@ -0,0 +1,26 @@
+use std::error::Error;
+use url::Url;
+
+fn try_parse_url(url: &str) -> Result<Url, Box<dyn Error>> {
+    Url::parse(url).map_err(|e| e.into())
+}
+
+pub fn try_parse_url_path(url: &str) -> Option<String> {
+    let full_url = if url.starts_with("/") {
+        format!("http://localhost{}", url)
+    } else {
+        url.to_string()
+    };
+
+    match try_parse_url(&full_url) {
+        Ok(parsed) => {
+            let path = parsed.path();
+            if path.is_empty() {
+                Some("/".to_string())
+            } else {
+                Some(path.to_string())
+            }
+        }
+        Err(_) => None,
+    }
+}

src/helpers/try_parse_url_path_test.rs

@@ -0,0 +1,83 @@
+#[cfg(test)]
+mod tests {
+    use crate::helpers::try_parse_url_path::try_parse_url_path;
+
+    #[test]
+    fn test_try_parse_url_path_nothing_found() {
+        assert_eq!(try_parse_url_path("abc"), None);
+    }
+
+    #[test]
+    fn test_try_parse_url_path_for_root() {
+        assert_eq!(try_parse_url_path("/"), Some("/".to_string()));
+    }
+
+    #[test]
+    fn test_try_parse_url_path_for_relative_url() {
+        assert_eq!(try_parse_url_path("/posts"), Some("/posts".to_string()));
+    }
+
+    #[test]
+    fn test_try_parse_url_path_for_relative_url_with_query() {
+        assert_eq!(
+            try_parse_url_path("/posts?abc=def"),
+            Some("/posts".to_string())
+        );
+    }
+
+    #[test]
+    fn test_try_parse_url_path_for_absolute_url() {
+        assert_eq!(
+            try_parse_url_path("http://localhost/posts/3"),
+            Some("/posts/3".to_string())
+        );
+    }
+
+    #[test]
+    fn test_try_parse_url_path_for_absolute_url_with_query() {
+        assert_eq!(
+            try_parse_url_path("http://localhost/posts/3?abc=def"),
+            Some("/posts/3".to_string())
+        );
+    }
+
+    #[test]
+    fn test_try_parse_url_path_for_absolute_url_with_hash() {
+        assert_eq!(
+            try_parse_url_path("http://localhost/posts/3#abc"),
+            Some("/posts/3".to_string())
+        );
+    }
+
+    #[test]
+    fn test_try_parse_url_path_for_absolute_url_with_query_and_hash() {
+        assert_eq!(
+            try_parse_url_path("http://localhost/posts/3?abc=def#ghi"),
+            Some("/posts/3".to_string())
+        );
+    }
+
+    #[test]
+    fn test_try_parse_url_path_for_absolute_url_with_query_and_hash_no_path() {
+        assert_eq!(
+            try_parse_url_path("http://localhost/?abc=def#ghi"),
+            Some("/".to_string())
+        );
+    }
+
+    #[test]
+    fn test_try_parse_url_path_for_absolute_url_with_query_no_path() {
+        assert_eq!(
+            try_parse_url_path("http://localhost?abc=def"),
+            Some("/".to_string())
+        );
+    }
+
+    #[test]
+    fn test_try_parse_url_path_for_absolute_url_with_hash_no_path() {
+        assert_eq!(
+            try_parse_url_path("http://localhost#abc"),
+            Some("/".to_string())
+        );
+    }
+}

src/lib.rs

@@ -2,4 +2,6 @@
  * Using FFI. Currently we support the following algorithms :
  * - Shell Injection (WIP)
  */
+pub mod helpers;
+pub mod route_builder;
 pub mod shell_injection;

src/route_builder/build_route_from_url.rs

@@ -0,0 +1,41 @@
+use crate::helpers::try_parse_url_path::try_parse_url_path;
+use crate::route_builder::replace_segment_with_param::replace_segment_with_param;
+
+use std::ffi::{CStr, CString};
+use std::os::raw::c_char;
+use std::str;
+
+#[no_mangle]
+pub extern "C" fn build_route_from_url(url: *const c_char) -> *mut c_char {
+    // Check if the pointers are null
+    if url.is_null() {
+        eprintln!("Received null pointer for URL.");
+        return CString::new("").expect("").into_raw();
+    }
+    let url_bytes = unsafe { CStr::from_ptr(url).to_bytes() };
+    let url_str = str::from_utf8(url_bytes).unwrap();
+
+    let route = build_route_from_url_str(url_str).unwrap_or("".to_string());
+    let c_string = CString::new(route).expect("CString::new failed");
+    c_string.into_raw()
+}
+
+pub fn build_route_from_url_str(url: &str) -> Option<String> {
+    let path = try_parse_url_path(url)?;
+
+    let route: String = path
+        .split('/')
+        .map(replace_segment_with_param)
+        .collect::<Vec<String>>()
+        .join("/");
+
+    if route == "/" {
+        return Some("/".to_string());
+    }
+
+    if route.ends_with('/') {
+        return Some(route[..route.len() - 1].to_string());
+    }
+
+    Some(route)
+}

src/route_builder/build_route_from_url_test.rs

@@ -0,0 +1,162 @@
+#[cfg(test)]
+mod tests {
+    use crate::route_builder::build_route_from_url::build_route_from_url_str;
+
+    #[test]
+    fn test_invalid_urls() {
+        assert_eq!(build_route_from_url_str(""), None);
+        assert_eq!(build_route_from_url_str("http"), None);
+    }
+
+    #[test]
+    fn test_root_urls() {
+        assert_eq!(build_route_from_url_str("/"), Some("/".to_string()));
+        assert_eq!(
+            build_route_from_url_str("http://localhost/"),
+            Some("/".to_string())
+        );
+    }
+
+    #[test]
+    fn test_replace_numbers() {
+        assert_eq!(
+            build_route_from_url_str("/posts/3"),
+            Some("/posts/:number".to_string())
+        );
+        assert_eq!(
+            build_route_from_url_str("http://localhost/posts/3"),
+            Some("/posts/:number".to_string())
+        );
+        assert_eq!(
+            build_route_from_url_str("http://localhost/posts/3/"),
+            Some("/posts/:number".to_string())
+        );
+        assert_eq!(
+            build_route_from_url_str("http://localhost/posts/3/comments/10"),
+            Some("/posts/:number/comments/:number".to_string())
+        );
+        assert_eq!(
+            build_route_from_url_str("/blog/2023/05/great-article"),
+            Some("/blog/:number/:number/great-article".to_string())
+        );
+    }
+
+    #[test]
+    fn test_replace_dates() {
+        assert_eq!(
+            build_route_from_url_str("/posts/2023-05-01"),
+            Some("/posts/:date".to_string())
+        );
+        assert_eq!(
+            build_route_from_url_str("/posts/2023-05-01/"),
+            Some("/posts/:date".to_string())
+        );
+        assert_eq!(
+            build_route_from_url_str("/posts/2023-05-01/comments/2023-05-01"),
+            Some("/posts/:date/comments/:date".to_string())
+        );
+        assert_eq!(
+            build_route_from_url_str("/posts/01-05-2023"),
+            Some("/posts/:date".to_string())
+        );
+    }
+
+    #[test]
+    fn test_ignore_comma_numbers() {
+        assert_eq!(
+            build_route_from_url_str("/posts/3,000"),
+            Some("/posts/3,000".to_string())
+        );
+    }
+
+    #[test]
+    fn test_ignore_api_version_numbers() {
+        assert_eq!(
+            build_route_from_url_str("/v1/posts/3"),
+            Some("/v1/posts/:number".to_string())
+        );
+    }
+
+    #[test]
+    fn test_replace_uuids() {
+        let uuids = [
+            "d9428888-122b-11e1-b85c-61cd3cbb3210",
+            "000003e8-2363-21ef-b200-325096b39f47",
+            "a981a0c2-68b1-35dc-bcfc-296e52ab01ec",
+            "109156be-c4fb-41ea-b1b4-efe1671c5836",
+            "90123e1c-7512-523e-bb28-76fab9f2f73d",
+            "1ef21d2f-1207-6660-8c4f-419efbd44d48",
+            "017f22e2-79b0-7cc3-98c4-dc0c0c07398f",
+            "0d8f23a0-697f-83ae-802e-48f3756dd581",
+        ];
+        for uuid in &uuids {
+            assert_eq!(
+                build_route_from_url_str(&format!("/posts/{}", uuid)),
+                Some("/posts/:uuid".to_string())
+            );
+        }
+    }
+
+    #[test]
+    fn test_ignore_invalid_uuids() {
+        assert_eq!(
+            build_route_from_url_str("/posts/00000000-0000-1000-6000-000000000000"),
+            Some("/posts/00000000-0000-1000-6000-000000000000".to_string())
+        );
+    }
+
+    #[test]
+    fn test_ignore_strings() {
+        assert_eq!(
+            build_route_from_url_str("/posts/abc"),
+            Some("/posts/abc".to_string())
+        );
+    }
+
+    #[test]
+    fn test_replace_email_addresses() {
+        assert_eq!(
+            build_route_from_url_str("/login/[email protected]"),
+            Some("/login/:email".to_string())
+        );
+        assert_eq!(
+            build_route_from_url_str("/login/[email protected]"),
+            Some("/login/:email".to_string())
+        );
+    }
+
+    #[test]
+    fn test_replace_ip_addresses() {
+        assert_eq!(
+            build_route_from_url_str("/block/1.2.3.4"),
+            Some("/block/:ip".to_string())
+        );
+        assert_eq!(
+            build_route_from_url_str("/block/2001:2:ffff:ffff:ffff:ffff:ffff:ffff"),
+            Some("/block/:ip".to_string())
+        );
+        assert_eq!(
+            build_route_from_url_str("/block/64:ff9a::255.255.255.255"),
+            Some("/block/:ip".to_string())
+        );
+        assert_eq!(
+            build_route_from_url_str("/block/100::"),
+            Some("/block/:ip".to_string())
+        );
+        assert_eq!(
+            build_route_from_url_str("/block/fec0::"),
+            Some("/block/:ip".to_string())
+        );
+        assert_eq!(
+            build_route_from_url_str("/block/227.202.96.196"),
+            Some("/block/:ip".to_string())
+        );
+    }
+    #[test]
+    fn test_replace_secrets() {
+        assert_eq!(
+            build_route_from_url_str("/confirm/CnJ4DunhYfv2db6T1FRfciRBHtlNKOYrjoz"),
+            Some("/confirm/:secret".to_string())
+        );
+    }
+}