whois_cli/
query.rs

1use std::io::{Read, Write};
2use std::net::TcpStream;
3use std::time::Duration;
4use anyhow::{Context, Result};
5use crate::servers::{WhoisServer, ServerSelector, DEFAULT_WHOIS_SERVER};
6use crate::protocol::WhoisColorProtocol;
7
8const TIMEOUT_SECONDS: u64 = 10;
9
10/// Check if a WHOIS response is effectively empty or indicates no results
11fn is_empty_result(response: &str) -> bool {
12    let response = response.trim();
13    
14    // Obviously empty
15    if response.is_empty() {
16        return true;
17    }
18    
19    // Common empty response indicators (case-insensitive)
20    let response_lower = response.to_lowercase();
21    let empty_indicators = [
22        "no found",
23        "no match",
24        "not found",
25        "no data found",
26        "no entries found",
27        "no records found",
28        "no such domain",
29        "no whois server is known",
30        "object does not exist",
31        "%error: no objects found",
32        "% no objects found",
33    ];
34    
35    for indicator in &empty_indicators {
36        if response_lower.contains(indicator) {
37            return true;
38        }
39    }
40    
41    // Check if response only contains comment lines (lines starting with % or #)
42    let content_lines: Vec<&str> = response
43        .lines()
44        .map(|line| line.trim())
45        .filter(|line| !line.is_empty())
46        .filter(|line| !line.starts_with('%') && !line.starts_with('#'))
47        .collect();
48    
49    if content_lines.is_empty() {
50        return true;
51    }
52    
53    // Check if response is very short (less than 30 characters) and likely just headers/boilerplate
54    // Only apply this for extremely short responses that have minimal content
55    if response.len() < 30 && content_lines.join(" ").len() < 10 {
56        return true;
57    }
58    
59    false
60}
61
62#[derive(Debug)]
63pub struct QueryResult {
64    pub response: String,
65    pub server_used: WhoisServer,
66    pub server_colored: bool,
67}
68
69impl QueryResult {
70    pub fn new(response: String, server_used: WhoisServer) -> Self {
71        Self { 
72            response, 
73            server_used,
74            server_colored: false,
75        }
76    }
77
78    pub fn new_with_color(response: String, server_used: WhoisServer, server_colored: bool) -> Self {
79        Self { 
80            response, 
81            server_used,
82            server_colored,
83        }
84    }
85}
86
87pub struct WhoisQuery {
88    verbose: bool,
89}
90
91impl WhoisQuery {
92    pub fn new(verbose: bool) -> Self {
93        Self { verbose }
94    }
95
96    /// Perform a direct WHOIS query to a specific server
97    pub fn query_direct(&self, query: &str, server: &WhoisServer) -> Result<String> {
98        let address = server.address();
99        
100        if self.verbose {
101            println!("Connecting to: {}", address);
102        }
103
104        let mut stream = TcpStream::connect(&address)
105            .with_context(|| format!("Cannot connect to WHOIS server: {}", address))?;
106        
107        stream.set_read_timeout(Some(Duration::from_secs(TIMEOUT_SECONDS)))
108            .context("Failed to set read timeout")?;
109        
110        stream.set_write_timeout(Some(Duration::from_secs(TIMEOUT_SECONDS)))
111            .context("Failed to set write timeout")?;
112        
113        let query_string = format!("{}\r\n", query);
114        stream.write_all(query_string.as_bytes())
115            .context("Failed to write query to WHOIS server")?;
116        
117        let mut response = String::new();
118        stream.read_to_string(&mut response)
119            .context("Failed to read response from WHOIS server")?;
120        
121        Ok(response)
122    }
123
124    /// Perform a WHOIS query with IANA referral if needed
125    pub fn query_with_referral(&self, query: &str, initial_server: &WhoisServer) -> Result<QueryResult> {
126        if initial_server.name == "IANA" {
127            if self.verbose {
128                println!("Querying IANA at: {}", initial_server.address());
129            }
130
131            // First query IANA
132            let iana_response = self.query_direct(query, initial_server)?;
133            
134            // Extract the referral WHOIS server from IANA's response
135            let whois_server_host = ServerSelector::extract_whois_server(&iana_response)
136                .unwrap_or_else(|| DEFAULT_WHOIS_SERVER.to_string());
137            
138            let final_server = WhoisServer::custom(whois_server_host, initial_server.port);
139            
140            if self.verbose {
141                if final_server.host != DEFAULT_WHOIS_SERVER {
142                    println!("IANA referred to: {}", final_server.host);
143                } else {
144                    println!("No referral found, using default: {}", DEFAULT_WHOIS_SERVER);
145                }
146            }
147            
148            // Query the actual WHOIS server
149            let final_response = self.query_direct(query, &final_server)?;
150            
151            Ok(QueryResult::new(final_response, final_server))
152        } else {
153            // Direct query to specified server
154            if self.verbose {
155                println!("Using {} server: {}", initial_server.name, initial_server.address());
156            }
157
158            let response = self.query_direct(query, initial_server)?;
159            Ok(QueryResult::new(response, initial_server.clone()))
160        }
161    }
162
163    /// Main query method that handles all logic
164    pub fn query(
165        &self,
166        domain: &str,
167        use_dn42: bool,
168        use_bgptools: bool,
169        explicit_server: Option<&str>,
170        port: u16,
171    ) -> Result<QueryResult> {
172        let server = ServerSelector::select_server(
173            domain,
174            use_dn42,
175            use_bgptools,
176            explicit_server,
177            port,
178        );
179
180        let result = self.query_with_referral(domain, &server)?;
181        
182        // Check if result is empty and fallback to RADB if needed
183        // Only fallback if we're not already using a specific server (DN42, BGPtools, or explicit server)
184        if is_empty_result(&result.response) && 
185           !use_dn42 && !use_bgptools && explicit_server.is_none() && 
186           server.name != "RADB" {
187            
188            if self.verbose {
189                println!("Empty result from RIR servers, trying RADB fallback...");
190            }
191            
192            return self.try_radb_fallback(domain, false, false, false, None);
193        }
194        
195        Ok(result)
196    }
197
198    /// Query with enhanced protocol support (v1.1 with markdown and images)
199    pub fn query_with_enhanced_protocol(
200        &self,
201        domain: &str,
202        use_dn42: bool,
203        use_bgptools: bool,
204        use_server_color: bool,
205        enable_markdown: bool,
206        enable_images: bool,
207        explicit_server: Option<&str>,
208        port: u16,
209        preferred_color_scheme: Option<&str>,
210    ) -> Result<QueryResult> {
211        let server = ServerSelector::select_server(
212            domain,
213            use_dn42,
214            use_bgptools,
215            explicit_server,
216            port,
217        );
218
219        let result = if use_server_color || enable_markdown || enable_images {
220            self.query_with_enhanced_protocol_impl(domain, &server, preferred_color_scheme, enable_markdown, enable_images)?
221        } else {
222            self.query_with_referral(domain, &server)?
223        };
224
225        // Check if result is empty and fallback to RADB if needed
226        // Only fallback if we're not already using a specific server (DN42, BGPtools, or explicit server)
227        if is_empty_result(&result.response) && 
228           !use_dn42 && !use_bgptools && explicit_server.is_none() && 
229           server.name != "RADB" {
230            
231            if self.verbose {
232                println!("Empty result from RIR servers, trying RADB fallback...");
233            }
234            
235            return self.try_radb_fallback(domain, use_server_color, enable_markdown, enable_images, preferred_color_scheme);
236        }
237
238        Ok(result)
239    }
240
241    /// Legacy method for backward compatibility
242    /// Query with color protocol support
243    pub fn query_with_color_protocol(
244        &self,
245        domain: &str,
246        use_dn42: bool,
247        use_bgptools: bool,
248        use_server_color: bool,
249        explicit_server: Option<&str>,
250        port: u16,
251        preferred_color_scheme: Option<&str>,
252    ) -> Result<QueryResult> {
253        let server = ServerSelector::select_server(
254            domain,
255            use_dn42,
256            use_bgptools,
257            explicit_server,
258            port,
259        );
260
261        let result = if use_server_color {
262            self.query_with_enhanced_protocol_impl(domain, &server, preferred_color_scheme, false, false)?
263        } else {
264            self.query_with_referral(domain, &server)?
265        };
266
267        // Check if result is empty and fallback to RADB if needed
268        // Only fallback if we're not already using a specific server (DN42, BGPtools, or explicit server)
269        if is_empty_result(&result.response) && 
270           !use_dn42 && !use_bgptools && explicit_server.is_none() && 
271           server.name != "RADB" {
272            
273            if self.verbose {
274                println!("Empty result from RIR servers, trying RADB fallback...");
275            }
276            
277            return self.try_radb_fallback(domain, use_server_color, false, false, preferred_color_scheme);
278        }
279
280        Ok(result)
281    }
282
283    /// Implementation of enhanced protocol query (v1.1)
284    fn query_with_enhanced_protocol_impl(
285        &self,
286        domain: &str,
287        server: &WhoisServer,
288        preferred_color_scheme: Option<&str>,
289        enable_markdown: bool,
290        enable_images: bool,
291    ) -> Result<QueryResult> {
292        let protocol = WhoisColorProtocol;
293        
294        if server.name == "IANA" {
295            // Handle IANA referral first
296            if self.verbose {
297                println!("Querying IANA at: {}", server.address());
298            }
299
300            let iana_response = self.query_direct(domain, server)?;
301            let whois_server_host = ServerSelector::extract_whois_server(&iana_response)
302                .unwrap_or_else(|| DEFAULT_WHOIS_SERVER.to_string());
303            
304            let final_server = WhoisServer::custom(whois_server_host, server.port);
305            
306            if self.verbose {
307                if final_server.host != DEFAULT_WHOIS_SERVER {
308                    println!("IANA referred to: {}", final_server.host);
309                } else {
310                    println!("No referral found, using default: {}", DEFAULT_WHOIS_SERVER);
311                }
312            }
313
314            // Try enhanced protocol with final server
315            return self.try_enhanced_protocol_query(domain, &final_server, &protocol, preferred_color_scheme, enable_markdown, enable_images);
316        } else {
317            // Direct server query with enhanced protocol
318            return self.try_enhanced_protocol_query(domain, server, &protocol, preferred_color_scheme, enable_markdown, enable_images);
319        }
320    }
321
322
323    /// Try enhanced protocol query with all v1.1 features
324    fn try_enhanced_protocol_query(
325        &self,
326        domain: &str,
327        server: &WhoisServer,
328        protocol: &WhoisColorProtocol,
329        preferred_color_scheme: Option<&str>,
330        enable_markdown: bool,
331        enable_images: bool,
332    ) -> Result<QueryResult> {
333        // Probe server capabilities
334        let capabilities = protocol.probe_capabilities(&server.address(), self.verbose)
335            .unwrap_or_default(); // Use default (no support) if probe fails
336
337        // Perform query based on capabilities
338        let response = protocol.query_with_enhanced_protocol(
339            &server.address(),
340            domain,
341            &capabilities,
342            preferred_color_scheme,
343            enable_markdown,
344            enable_images,
345            self.verbose
346        )?;
347
348        let server_colored = protocol.is_server_colored(&response);
349        Ok(QueryResult::new_with_color(response, server.clone(), server_colored))
350    }
351
352    /// Try RADB fallback when RIR servers return empty results
353    fn try_radb_fallback(
354        &self,
355        domain: &str,
356        use_server_color: bool,
357        enable_markdown: bool,
358        enable_images: bool,
359        preferred_color_scheme: Option<&str>,
360    ) -> Result<QueryResult> {
361        let radb_server = WhoisServer::radb();
362        
363        if self.verbose {
364            println!("Querying RADB at: {}", radb_server.address());
365        }
366        
367        if use_server_color || enable_markdown || enable_images {
368            // Try enhanced protocol with RADB
369            self.query_with_enhanced_protocol_impl(domain, &radb_server, preferred_color_scheme, enable_markdown, enable_images)
370        } else {
371            // Direct query to RADB
372            let response = self.query_direct(domain, &radb_server)?;
373            Ok(QueryResult::new(response, radb_server))
374        }
375    }
376
377}
378
379#[cfg(test)]
380mod tests {
381    use super::*;
382
383    #[test]
384    fn test_is_empty_result_completely_empty() {
385        assert!(is_empty_result(""));
386        assert!(is_empty_result("   "));
387        assert!(is_empty_result("\n\n\n"));
388    }
389
390    #[test]
391    fn test_is_empty_result_common_indicators() {
392        assert!(is_empty_result("No Found"));
393        assert!(is_empty_result("NO MATCH"));
394        assert!(is_empty_result("not found"));
395        assert!(is_empty_result("No data found"));
396        assert!(is_empty_result("No entries found"));
397        assert!(is_empty_result("No records found"));
398        assert!(is_empty_result("No such domain"));
399        assert!(is_empty_result("No whois server is known"));
400        assert!(is_empty_result("Object does not exist"));
401        assert!(is_empty_result("%Error: No objects found"));
402        assert!(is_empty_result("% No objects found"));
403    }
404
405    #[test]
406    fn test_is_empty_result_short_responses() {
407        assert!(is_empty_result("Short"));
408        assert!(is_empty_result("Tiny"));
409        assert!(!is_empty_result("Very short response")); // This is now long enough to be considered valid
410        assert!(!is_empty_result("This is a longer response that should be considered valid content with enough information"));
411    }
412
413    #[test]
414    fn test_is_empty_result_comment_only() {
415        assert!(is_empty_result("% Comment only\n% Another comment\n# More comments"));
416        assert!(is_empty_result("% This is just comments\n\n% More comments"));
417        assert!(!is_empty_result("% Comment\nactual content\n% Another comment"));
418        assert!(!is_empty_result("Some real content\n% with comment"));
419    }
420
421    #[test]
422    fn test_is_empty_result_valid_content() {
423        let valid_content = r#"
424domain:         example.com
425descr:          Example Domain
426admin-c:        ADMIN123
427tech-c:         TECH456
428status:         ASSIGNED
429mnt-by:         EXAMPLE-MNT
430created:        2020-01-01T00:00:00Z
431last-modified:  2020-12-31T23:59:59Z
432source:         RIPE
433        "#;
434        assert!(!is_empty_result(valid_content));
435    }
436
437    #[test]
438    fn test_radb_server_creation() {
439        let radb = WhoisServer::radb();
440        assert_eq!(radb.host, "whois.radb.net");
441        assert_eq!(radb.port, 43);
442        assert_eq!(radb.name, "RADB");
443        assert_eq!(radb.address(), "whois.radb.net:43");
444    }
445}