From 268ae337fd86d108009c8ed2bce6c01078b577e5 Mon Sep 17 00:00:00 2001
From: camo <196579573+camo-rs@users.noreply.github.com>
Date: Wed, 5 Mar 2025 12:59:36 -0500
Subject: [PATCH] show client names, show DNS rewrites

---
 src/fetch/fetch_query_log.rs |  7 ++++++
 src/fetch/fetch_stats.rs     | 43 ++++++++++++++++++++++++++++++------
 src/main.rs                  |  9 +++++++-
 src/widgets/table.rs         | 12 +++++++++-
 4 files changed, 62 insertions(+), 9 deletions(-)

diff --git a/src/fetch/fetch_query_log.rs b/src/fetch/fetch_query_log.rs
index f0de9bf..ccbb1e3 100644
--- a/src/fetch/fetch_query_log.rs
+++ b/src/fetch/fetch_query_log.rs
@@ -16,6 +16,8 @@ pub struct Query {
     #[serde(rename = "elapsedMs")]
     pub elapsed_ms: String,
     pub question: Question,
+    #[serde(default)]
+    pub client_info: Option<ClientInfo>,
     pub reason: String,
     pub time: String,
 }
@@ -28,6 +30,11 @@ pub struct Question {
     pub question_type: String,
 }
 
+#[derive(Deserialize)]
+pub struct ClientInfo {
+    pub name: Option<String>,
+}
+
 pub async fn fetch_adguard_query_log(
   client: &reqwest::Client,
   endpoint: &str,
diff --git a/src/fetch/fetch_stats.rs b/src/fetch/fetch_stats.rs
index 74516ad..401795c 100644
--- a/src/fetch/fetch_stats.rs
+++ b/src/fetch/fetch_stats.rs
@@ -14,6 +14,17 @@ pub struct DomainData {
     pub count: i32,
 }
 
+#[derive(Debug, Deserialize, Clone)]
+pub struct ClientsResponse {
+    clients: Vec<ClientEntry>,
+}
+
+#[derive(Debug, Deserialize, Clone)]
+pub struct ClientEntry {
+    name: Option<String>,
+    ids: Vec<String>,
+}
+
 #[derive(Debug, Deserialize, Clone)]
 pub struct StatsResponse {
     pub num_dns_queries: u64,
@@ -45,21 +56,39 @@ pub async fn fetch_adguard_stats(
     endpoint: &str,
     username: &str,
     password: &str,
-) -> Result<StatsResponse, anyhow::Error> {
+) -> Result<(StatsResponse, HashMap<String, String>), anyhow::Error> {
     let auth_string = format!("{}:{}", username, password);
     let auth_header_value = format!("Basic {}", base64::encode(&auth_string));
     let mut headers = reqwest::header::HeaderMap::new();
     headers.insert(AUTHORIZATION, auth_header_value.parse()?);
     headers.insert(CONTENT_LENGTH, HeaderValue::from_static("0"));
 
-    let url = format!("{}/control/stats", endpoint);
-    let response = client.get(&url).headers(headers).send().await?;
-    if !response.status().is_success() {
-        return Err(anyhow::anyhow!("Request failed with status code {}", response.status()));
+    let clients_url = format!("{}/control/clients", endpoint);
+    let clients_response = client.get(&clients_url).headers(headers.clone()).send().await?;
+    if !clients_response.status().is_success() {
+        return Err(anyhow::anyhow!("Clients request failed with status code {}", clients_response.status()));
+    }
+
+    let possible_clients: Option<ClientsResponse> = clients_response.json().await.unwrap_or(None);
+    let mut client_map = HashMap::new();
+    if let Some(clients) = possible_clients {
+        for client in clients.clients {
+            if let Some(name) = client.name {
+                for ip in client.ids {
+                    client_map.insert(ip, name.clone());
+                }
+            }
+        }
+    }
+
+    let stats_url = format!("{}/control/stats", endpoint);
+    let stats_response = client.get(&stats_url).headers(headers).send().await?;
+    if !stats_response.status().is_success() {
+        return Err(anyhow::anyhow!("Stats request failed with status code {}", stats_response.status()));
     }
 
-    let data = response.json().await?;
-    Ok(data)
+    let stats: StatsResponse = stats_response.json().await?;
+    Ok((stats, client_map))
 }
 
 /// Deserialize a list of domains from the JSON data
diff --git a/src/main.rs b/src/main.rs
index ccc2caf..968badc 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -60,7 +60,14 @@ async fn run() -> anyhow::Result<()> {
                     return Err(anyhow::anyhow!("Failed to send query data"));
                 }
                 
-                let stats = fetch_adguard_stats(&client, &hostname, &username, &password).await?;
+                let (mut stats, client_map) = fetch_adguard_stats(&client, &hostname, &username, &password).await?;
+                
+                for domain in &mut stats.top_clients {
+                    if let Some(client_name) = client_map.get(&domain.name) {
+                        domain.name = client_name.clone();
+                    }
+                }
+                
                 if stats_tx.send(stats).await.is_err() {
                     return Err(anyhow::anyhow!("Failed to send stats data"));
                 }
diff --git a/src/widgets/table.rs b/src/widgets/table.rs
index 57cbbcb..590799f 100644
--- a/src/widgets/table.rs
+++ b/src/widgets/table.rs
@@ -17,7 +17,13 @@ pub fn make_query_table(data: &[Query], width: u16) -> Table<'_> {
       let question = Cell::from(make_request_cell(&query.question).unwrap())
           .style(Style::default().add_modifier(Modifier::BOLD));
 
-      let client = Cell::from(query.client.as_str())
+      let client_names = query.client_info
+          .as_ref()
+          .and_then(|info| info.name.as_deref())
+          .filter(|name| !name.is_empty())
+          .map_or_else(|| query.client.as_str(), |name| name);
+      
+      let client = Cell::from(client_names)
           .style(Style::default().fg(Color::Blue));
 
       let (time_taken, elapsed_color) = make_time_taken_and_color(&query.elapsed_ms).unwrap();
@@ -125,6 +131,8 @@ fn make_row_color(reason: &str) -> Color {
       Color::Green
   } else if reason == "FilteredBlackList" {
       Color::Red
+  } else if reason == "Rewrite" {
+      Color::LightGreen
   } else {
       Color::Yellow
   }
@@ -137,6 +145,8 @@ fn block_status_text(reason: &str, cached: bool) -> (String, Color) {
       ("Blacklisted".to_string(), Color::Red)
   } else if cached {
       ("Cached".to_string(), Color::Cyan)
+  } else if reason == "Rewrite" {
+      ("Rewrite".to_string(), Color::LightGreen)
   } else if reason == "NotFilteredNotFound" {
       ("Allowed".to_string(), Color::Green)
   } else {