diff --git a/src/utils/api.rs b/src/utils/api.rs index 7bec827..bf7f6b3 100644 --- a/src/utils/api.rs +++ b/src/utils/api.rs @@ -15,6 +15,10 @@ use crate::log::debug; const CHUNK_SIZE: usize = 50 * 1024 * 1024; // 50 MB const API_BASE: &str = "/api/v1"; +fn get_source() -> String { + std::env::var("CORGEA_SOURCE").unwrap_or_else(|_| "cli".to_string()) +} + pub fn http_client() -> reqwest::blocking::Client { let mut builder = reqwest::blocking::Client::builder().timeout(std::time::Duration::from_secs(5 * 30)); @@ -76,6 +80,7 @@ pub fn upload_zip( let response_object = client .post(format!("{}{}/start-scan", url, API_BASE)) .header("CORGEA-TOKEN", token) + .header("CORGEA-SOURCE", get_source()) .query(&[ ("scan_type", "blast"), ]) @@ -165,6 +170,7 @@ pub fn upload_zip( let response = match client .patch(format!("{}{}/start-scan/{}/", url, API_BASE, transfer_id)) .header("CORGEA-TOKEN", token) + .header("CORGEA-SOURCE", get_source()) .header("Upload-Offset", offset.to_string()) .header("Upload-Length", file_size.to_string()) .header("Upload-Name", file_name) @@ -281,6 +287,7 @@ pub fn get_scan_issues( let client = http_client(); let mut headers = HeaderMap::new(); headers.insert("CORGEA-TOKEN", token.parse().unwrap()); + headers.insert("CORGEA-SOURCE", get_source().parse().unwrap()); debug(&format!("Sending request to URL: {}", url)); debug(&format!("Headers: {:?}", headers)); @@ -314,6 +321,7 @@ pub fn get_scan(url: &str, token: &str, scan_id: &str) -> Result Result Result