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
10fn is_empty_result(response: &str) -> bool {
12 let response = response.trim();
13
14 if response.is_empty() {
16 return true;
17 }
18
19 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 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 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 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 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 let iana_response = self.query_direct(query, initial_server)?;
133
134 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 let final_response = self.query_direct(query, &final_server)?;
150
151 Ok(QueryResult::new(final_response, final_server))
152 } else {
153 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 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 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 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 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 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 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 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 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 return self.try_enhanced_protocol_query(domain, &final_server, &protocol, preferred_color_scheme, enable_markdown, enable_images);
316 } else {
317 return self.try_enhanced_protocol_query(domain, server, &protocol, preferred_color_scheme, enable_markdown, enable_images);
319 }
320 }
321
322
323 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 let capabilities = protocol.probe_capabilities(&server.address(), self.verbose)
335 .unwrap_or_default(); 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 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 self.query_with_enhanced_protocol_impl(domain, &radb_server, preferred_color_scheme, enable_markdown, enable_images)
370 } else {
371 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")); 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}