whois_cli/
protocol.rs

1use std::io::{Read, Write};
2use std::net::TcpStream;
3use std::time::Duration;
4use anyhow::{Context, Result};
5
6/// WHOIS-COLOR Protocol v1.1
7/// A backward-compatible extension protocol for server-side colorization,
8/// Markdown rendering, and image display support
9pub struct WhoisColorProtocol;
10
11/// Server capability information
12#[derive(Debug, Clone, PartialEq)]
13pub struct ServerCapabilities {
14    pub supports_color: bool,
15    pub color_schemes: Vec<String>,
16    pub protocol_version: String,
17    pub supports_markdown: bool,
18    pub supports_images: bool,
19    pub image_formats: Vec<String>,
20}
21
22impl Default for ServerCapabilities {
23    fn default() -> Self {
24        Self {
25            supports_color: false,
26            color_schemes: vec![],
27            protocol_version: "none".to_string(),
28            supports_markdown: false,
29            supports_images: false,
30            image_formats: vec![],
31        }
32    }
33}
34
35/// Protocol constants
36pub const PROTOCOL_VERSION: &str = "1.1";
37pub const LEGACY_VERSION: &str = "1.0";
38pub const CAPABILITY_PROBE: &str = "X-WHOIS-COLOR-PROBE: v1.1\r\n";
39pub const COLOR_REQUEST_PREFIX: &str = "X-WHOIS-COLOR: ";
40pub const MARKDOWN_REQUEST_PREFIX: &str = "X-WHOIS-MARKDOWN: ";
41pub const IMAGE_REQUEST_PREFIX: &str = "X-WHOIS-IMAGES: ";
42pub const CAPABILITY_RESPONSE_PREFIX: &str = "X-WHOIS-COLOR-SUPPORT: ";
43pub const CAPABILITY_TIMEOUT_MS: u64 = 2000; // 2 seconds for capability probe
44
45impl WhoisColorProtocol {
46    /// Probe server for color protocol support
47    /// This method sends a capability probe and waits for a response
48    /// If no response or timeout, assumes standard WHOIS server
49    pub fn probe_capabilities(
50        &self, 
51        server_address: &str,
52        verbose: bool
53    ) -> Result<ServerCapabilities> {
54        if verbose {
55            println!("Probing color capabilities for: {}", server_address);
56        }
57
58        let mut stream = TcpStream::connect(server_address)
59            .with_context(|| format!("Cannot connect to server for capability probe: {}", server_address))?;
60        
61        // Set shorter timeout for capability probe
62        stream.set_read_timeout(Some(Duration::from_millis(CAPABILITY_TIMEOUT_MS)))
63            .context("Failed to set read timeout for capability probe")?;
64        
65        stream.set_write_timeout(Some(Duration::from_millis(CAPABILITY_TIMEOUT_MS)))
66            .context("Failed to set write timeout for capability probe")?;
67
68        // Send capability probe
69        // Format: "X-WHOIS-COLOR-PROBE: v1.0\r\n\r\n"
70        let probe_query = format!("{}\r\n", CAPABILITY_PROBE);
71        
72        if let Err(_) = stream.write_all(probe_query.as_bytes()) {
73            // If write fails, assume standard WHOIS server
74            if verbose {
75                println!("Capability probe write failed, assuming standard WHOIS");
76            }
77            return Ok(ServerCapabilities::default());
78        }
79
80        // Try to read response
81        let mut response = String::new();
82        match stream.read_to_string(&mut response) {
83            Ok(_) => {
84                let capabilities = self.parse_capability_response(&response);
85                if verbose {
86                    println!("Server capabilities: {:?}", capabilities);
87                }
88                Ok(capabilities)
89            }
90            Err(_) => {
91                // Timeout or read error - assume standard WHOIS server
92                if verbose {
93                    println!("No capability response, assuming standard WHOIS");
94                }
95                Ok(ServerCapabilities::default())
96            }
97        }
98    }
99
100    /// Parse capability response from server
101    /// Expected format: "X-WHOIS-COLOR-SUPPORT: v1.1 schemes=ripe,bgptools,mtf markdown=true images=png,jpg\r\n"
102    /// Legacy format: "X-WHOIS-COLOR-SUPPORT: v1.0 schemes=ripe,bgptools,mtf\r\n"
103    fn parse_capability_response(&self, response: &str) -> ServerCapabilities {
104        for line in response.lines() {
105            let line = line.trim();
106            if line.starts_with(CAPABILITY_RESPONSE_PREFIX) {
107                return self.parse_capability_line(&line[CAPABILITY_RESPONSE_PREFIX.len()..]);
108            }
109        }
110        ServerCapabilities::default()
111    }
112
113    /// Parse a single capability line
114    /// Format v1.1: "v1.1 schemes=ripe,bgptools,mtf markdown=true images=png,jpg"
115    /// Format v1.0: "v1.0 schemes=ripe,bgptools,mtf"
116    fn parse_capability_line(&self, capability_data: &str) -> ServerCapabilities {
117        let parts: Vec<&str> = capability_data.split_whitespace().collect();
118        if parts.is_empty() {
119            return ServerCapabilities::default();
120        }
121
122        let protocol_version = parts[0].to_string();
123        let mut capabilities = ServerCapabilities {
124            supports_color: true,
125            protocol_version: protocol_version.clone(),
126            color_schemes: vec![],
127            supports_markdown: false,
128            supports_images: false,
129            image_formats: vec![],
130        };
131
132        // Parse additional parameters
133        for part in &parts[1..] {
134            if let Some(schemes_part) = part.strip_prefix("schemes=") {
135                capabilities.color_schemes = schemes_part
136                    .split(',')
137                    .map(|s| s.trim().to_string())
138                    .filter(|s| !s.is_empty())
139                    .collect();
140            } else if let Some(markdown_part) = part.strip_prefix("markdown=") {
141                capabilities.supports_markdown = markdown_part == "true";
142            } else if let Some(images_part) = part.strip_prefix("images=") {
143                capabilities.supports_images = !images_part.is_empty();
144                capabilities.image_formats = images_part
145                    .split(',')
146                    .map(|s| s.trim().to_string())
147                    .filter(|s| !s.is_empty())
148                    .collect();
149            }
150        }
151
152        capabilities
153    }
154
155    /// Perform query with enhanced protocol support (color, markdown, images)
156    /// Falls back gracefully for older servers
157    pub fn query_with_enhanced_protocol(
158        &self,
159        server_address: &str,
160        query: &str,
161        capabilities: &ServerCapabilities,
162        preferred_scheme: Option<&str>,
163        enable_markdown: bool,
164        enable_images: bool,
165        verbose: bool
166    ) -> Result<String> {
167        let mut stream = TcpStream::connect(server_address)
168            .with_context(|| format!("Cannot connect to WHOIS server: {}", server_address))?;
169        
170        stream.set_read_timeout(Some(Duration::from_secs(10)))
171            .context("Failed to set read timeout")?;
172        
173        stream.set_write_timeout(Some(Duration::from_secs(10)))
174            .context("Failed to set write timeout")?;
175
176        let query_string = if capabilities.supports_color || capabilities.supports_markdown || capabilities.supports_images {
177            self.build_enhanced_query(query, capabilities, preferred_scheme, enable_markdown, enable_images, verbose)
178        } else {
179            // Standard WHOIS query
180            format!("{}\r\n", query)
181        };
182
183        if verbose {
184            if capabilities.supports_color {
185                println!("Sending color-enabled query");
186            }
187            if capabilities.supports_markdown && enable_markdown {
188                println!("Requesting Markdown format");
189            }
190            if capabilities.supports_images && enable_images {
191                println!("Requesting image support");
192            }
193        }
194
195        stream.write_all(query_string.as_bytes())
196            .context("Failed to write query to WHOIS server")?;
197        
198        let mut response = String::new();
199        stream.read_to_string(&mut response)
200            .context("Failed to read response from WHOIS server")?;
201        
202        Ok(response)
203    }
204
205    /// Build query string with enhanced protocol headers
206    /// Format v1.1: "X-WHOIS-COLOR: scheme=ripe\r\nX-WHOIS-MARKDOWN: true\r\nX-WHOIS-IMAGES: png,jpg\r\nquery\r\n"
207    /// Format v1.0: "X-WHOIS-COLOR: scheme=ripe\r\nquery\r\n"
208    fn build_enhanced_query(
209        &self,
210        query: &str,
211        capabilities: &ServerCapabilities,
212        preferred_scheme: Option<&str>,
213        enable_markdown: bool,
214        enable_images: bool,
215        verbose: bool
216    ) -> String {
217        let mut headers = String::new();
218        
219        // Add color header if supported
220        if capabilities.supports_color {
221            if let Some(scheme) = self.select_color_scheme(capabilities, preferred_scheme) {
222                if verbose {
223                    println!("Requesting server-side coloring with scheme: {}", scheme);
224                }
225                headers.push_str(&format!("{}scheme={}\r\n", COLOR_REQUEST_PREFIX, scheme));
226            }
227        }
228        
229        // Add markdown header if supported and requested
230        if capabilities.supports_markdown && enable_markdown {
231            headers.push_str(&format!("{}true\r\n", MARKDOWN_REQUEST_PREFIX));
232        }
233        
234        // Add images header if supported and requested
235        if capabilities.supports_images && enable_images && !capabilities.image_formats.is_empty() {
236            let formats = capabilities.image_formats.join(",");
237            headers.push_str(&format!("{}{}\r\n", IMAGE_REQUEST_PREFIX, formats));
238        }
239        
240        if headers.is_empty() {
241            // No protocol features, use standard query
242            format!("{}\r\n", query)
243        } else {
244            // Enhanced protocol query
245            format!("{}{}\r\n", headers, query)
246        }
247    }
248
249    /// Legacy method for backward compatibility
250    /// Build query string with color protocol headers
251    /// Format: "X-WHOIS-COLOR: scheme=ripe\r\nquery\r\n"
252    #[allow(dead_code)]
253    fn build_color_query(
254        &self,
255        query: &str,
256        capabilities: &ServerCapabilities,
257        preferred_scheme: Option<&str>,
258        verbose: bool
259    ) -> String {
260        let scheme = self.select_color_scheme(capabilities, preferred_scheme);
261        
262        if let Some(scheme) = scheme {
263            if verbose {
264                println!("Requesting server-side coloring with scheme: {}", scheme);
265            }
266            format!("{}scheme={}\r\n{}\r\n", COLOR_REQUEST_PREFIX, scheme, query)
267        } else {
268            // No suitable scheme, use standard query
269            if verbose {
270                println!("No suitable color scheme, falling back to standard query");
271            }
272            format!("{}\r\n", query)
273        }
274    }
275
276    /// Select appropriate color scheme based on server capabilities and preference
277    fn select_color_scheme(
278        &self,
279        capabilities: &ServerCapabilities,
280        preferred_scheme: Option<&str>
281    ) -> Option<String> {
282        if !capabilities.supports_color {
283            return None;
284        }
285
286        // If preferred scheme is supported, use it
287        if let Some(preferred) = preferred_scheme {
288            if capabilities.color_schemes.contains(&preferred.to_string()) {
289                return Some(preferred.to_string());
290            }
291        }
292
293        // Otherwise, use first available scheme
294        capabilities.color_schemes.first().cloned()
295    }
296
297    /// Check if response contains server-generated colors
298    /// Server-colored responses should contain color control sequences
299    pub fn is_server_colored(&self, response: &str) -> bool {
300        // Check for ANSI color escape sequences
301        response.contains("\x1b[") || 
302        // Check for color protocol markers (optional)
303        response.contains("X-WHOIS-COLOR-APPLIED:")
304    }
305}
306
307#[cfg(test)]
308mod tests {
309    use super::*;
310
311    #[test]
312    fn test_parse_capability_response_v10() {
313        let protocol = WhoisColorProtocol;
314        
315        let response = "X-WHOIS-COLOR-SUPPORT: v1.0 schemes=ripe,bgptools,mtf\r\n";
316        let capabilities = protocol.parse_capability_response(response);
317        
318        assert!(capabilities.supports_color);
319        assert_eq!(capabilities.protocol_version, "v1.0");
320        assert_eq!(capabilities.color_schemes, vec!["ripe", "bgptools", "mtf"]);
321        assert!(!capabilities.supports_markdown); // v1.0 doesn't support markdown
322        assert!(!capabilities.supports_images); // v1.0 doesn't support images
323    }
324
325    #[test]
326    fn test_parse_capability_response_v11() {
327        let protocol = WhoisColorProtocol;
328        
329        let response = "X-WHOIS-COLOR-SUPPORT: v1.1 schemes=ripe,bgptools,mtf markdown=true images=png,jpg\r\n";
330        let capabilities = protocol.parse_capability_response(response);
331        
332        assert!(capabilities.supports_color);
333        assert_eq!(capabilities.protocol_version, "v1.1");
334        assert_eq!(capabilities.color_schemes, vec!["ripe", "bgptools", "mtf"]);
335        assert!(capabilities.supports_markdown);
336        assert!(capabilities.supports_images);
337        assert_eq!(capabilities.image_formats, vec!["png", "jpg"]);
338    }
339
340    #[test]
341    fn test_parse_capability_response_minimal() {
342        let protocol = WhoisColorProtocol;
343        
344        let response = "X-WHOIS-COLOR-SUPPORT: v1.0\r\n";
345        let capabilities = protocol.parse_capability_response(response);
346        
347        assert!(capabilities.supports_color);
348        assert_eq!(capabilities.protocol_version, "v1.0");
349        assert!(capabilities.color_schemes.is_empty());
350    }
351
352    #[test]
353    fn test_parse_capability_response_no_support() {
354        let protocol = WhoisColorProtocol;
355        
356        let response = "Some other response\r\n";
357        let capabilities = protocol.parse_capability_response(response);
358        
359        assert!(!capabilities.supports_color);
360        assert_eq!(capabilities.protocol_version, "none");
361        assert!(capabilities.color_schemes.is_empty());
362    }
363
364    #[test]
365    fn test_select_color_scheme_preferred() {
366        let protocol = WhoisColorProtocol;
367        let capabilities = ServerCapabilities {
368            supports_color: true,
369            color_schemes: vec!["ripe".to_string(), "bgptools".to_string()],
370            protocol_version: "v1.0".to_string(),
371            supports_markdown: false,
372            supports_images: false,
373            image_formats: vec![],
374        };
375        
376        let scheme = protocol.select_color_scheme(&capabilities, Some("bgptools"));
377        assert_eq!(scheme, Some("bgptools".to_string()));
378    }
379
380    #[test]
381    fn test_select_color_scheme_fallback() {
382        let protocol = WhoisColorProtocol;
383        let capabilities = ServerCapabilities {
384            supports_color: true,
385            color_schemes: vec!["ripe".to_string(), "bgptools".to_string()],
386            protocol_version: "v1.0".to_string(),
387            supports_markdown: false,
388            supports_images: false,
389            image_formats: vec![],
390        };
391        
392        let scheme = protocol.select_color_scheme(&capabilities, Some("invalid"));
393        assert_eq!(scheme, Some("ripe".to_string()));
394    }
395
396    #[test]
397    fn test_select_color_scheme_no_support() {
398        let protocol = WhoisColorProtocol;
399        let capabilities = ServerCapabilities::default();
400        
401        let scheme = protocol.select_color_scheme(&capabilities, Some("ripe"));
402        assert_eq!(scheme, None);
403    }
404
405    #[test]
406    fn test_build_enhanced_query_v10_compat() {
407        let protocol = WhoisColorProtocol;
408        let capabilities = ServerCapabilities {
409            supports_color: true,
410            color_schemes: vec!["ripe".to_string()],
411            protocol_version: "v1.0".to_string(),
412            supports_markdown: false,
413            supports_images: false,
414            image_formats: vec![],
415        };
416        
417        let query = protocol.build_enhanced_query("example.com", &capabilities, Some("ripe"), false, false, false);
418        assert_eq!(query, "X-WHOIS-COLOR: scheme=ripe\r\nexample.com\r\n");
419    }
420
421    #[test]
422    fn test_build_enhanced_query_v11_full() {
423        let protocol = WhoisColorProtocol;
424        let capabilities = ServerCapabilities {
425            supports_color: true,
426            color_schemes: vec!["ripe".to_string()],
427            protocol_version: "v1.1".to_string(),
428            supports_markdown: true,
429            supports_images: true,
430            image_formats: vec!["png".to_string(), "jpg".to_string()],
431        };
432        
433        let query = protocol.build_enhanced_query("example.com", &capabilities, Some("ripe"), true, true, false);
434        let expected = "X-WHOIS-COLOR: scheme=ripe\r\nX-WHOIS-MARKDOWN: true\r\nX-WHOIS-IMAGES: png,jpg\r\nexample.com\r\n";
435        assert_eq!(query, expected);
436    }
437
438    #[test]
439    fn test_build_color_query_legacy() {
440        let protocol = WhoisColorProtocol;
441        let capabilities = ServerCapabilities {
442            supports_color: true,
443            color_schemes: vec!["ripe".to_string()],
444            protocol_version: "v1.0".to_string(),
445            supports_markdown: false,
446            supports_images: false,
447            image_formats: vec![],
448        };
449        
450        let query = protocol.build_color_query("example.com", &capabilities, Some("ripe"), false);
451        assert_eq!(query, "X-WHOIS-COLOR: scheme=ripe\r\nexample.com\r\n");
452    }
453
454    #[test]
455    fn test_build_color_query_no_scheme() {
456        let protocol = WhoisColorProtocol;
457        let capabilities = ServerCapabilities::default();
458        
459        let query = protocol.build_color_query("example.com", &capabilities, Some("ripe"), false);
460        assert_eq!(query, "example.com\r\n");
461    }
462
463    #[test]
464    fn test_is_server_colored() {
465        let protocol = WhoisColorProtocol;
466        
467        assert!(protocol.is_server_colored("text with \x1b[31mcolor\x1b[0m"));
468        assert!(protocol.is_server_colored("X-WHOIS-COLOR-APPLIED: ripe\ntext"));
469        assert!(!protocol.is_server_colored("plain text"));
470    }
471}
whois_cli/protocol.rs

whois_cli/
protocol.rs
