whois_rust/
who_is.rs

1use std::{
2    collections::HashMap,
3    fs::File,
4    io::{Read, Write},
5    net::{SocketAddr, TcpStream, ToSocketAddrs},
6    path::Path,
7    str::FromStr,
8    time::Duration,
9};
10
11use once_cell::sync::Lazy;
12use regex::Regex;
13use serde_json::{Map, Value};
14#[cfg(feature = "tokio")]
15use tokio::io::{AsyncReadExt, AsyncWriteExt};
16use trust_dns_client::{
17    client::{Client, SyncClient},
18    op::DnsResponse,
19    rr::{DNSClass, Name, RData, Record, RecordType},
20    udp::UdpClientConnection,
21};
22use validators::{models::Host, prelude::*};
23
24use crate::{WhoIsError, WhoIsLookupOptions, WhoIsServerValue};
25
26const DEFAULT_WHOIS_HOST_PORT: u16 = 43;
27const DEFAULT_WHOIS_HOST_QUERY: &str = "$addr\r\n";
28
29static RE_SERVER: Lazy<Regex> = Lazy::new(|| {
30    Regex::new(r"(ReferralServer|Registrar Whois|Whois Server|WHOIS Server|Registrar WHOIS Server):[^\S\n]*(r?whois://)?(.*)").unwrap()
31});
32
33/// The `WhoIs` structure stores the list of WHOIS servers in-memory.
34#[derive(Debug, Clone)]
35pub struct WhoIs {
36    map: HashMap<String, WhoIsServerValue>,
37    ip:  WhoIsServerValue,
38}
39
40impl WhoIs {
41    /// Create a `WhoIs` instance which doesn't have a WHOIS server list. You should provide the host that is used for query ip. You may want to use the host `"whois.arin.net"`.
42    pub fn from_host<T: AsRef<str>>(host: T) -> Result<WhoIs, WhoIsError> {
43        Ok(Self {
44            map: HashMap::new(), ip: WhoIsServerValue::from_string(host)?
45        })
46    }
47
48    /// Read the list of WHOIS servers (JSON data) from a file to create a `WhoIs` instance.
49    #[inline]
50    pub fn from_path<P: AsRef<Path>>(path: P) -> Result<WhoIs, WhoIsError> {
51        let path = path.as_ref();
52
53        let file = File::open(path)?;
54
55        let map: Map<String, Value> = serde_json::from_reader(file)?;
56
57        Self::from_inner(map)
58    }
59
60    #[cfg(feature = "tokio")]
61    /// Read the list of WHOIS servers (JSON data) from a file to create a `WhoIs` instance. For `serde_json` doesn't support async functions, consider just using the `from_path` function.
62    #[inline]
63    pub async fn from_path_async<P: AsRef<Path>>(path: P) -> Result<WhoIs, WhoIsError> {
64        let file = tokio::fs::read(path).await?;
65
66        let map: Map<String, Value> = serde_json::from_slice(file.as_slice())?;
67
68        Self::from_inner(map)
69    }
70
71    /// Read the list of WHOIS servers (JSON data) from a string to create a `WhoIs` instance.
72    #[inline]
73    pub fn from_string<S: AsRef<str>>(string: S) -> Result<WhoIs, WhoIsError> {
74        let string = string.as_ref();
75
76        let map: Map<String, Value> = serde_json::from_str(string)?;
77
78        Self::from_inner(map)
79    }
80
81    fn from_inner(mut map: Map<String, Value>) -> Result<WhoIs, WhoIsError> {
82        let ip = match map.remove("_") {
83            Some(server) => {
84                if let Value::Object(server) = server {
85                    match server.get("ip") {
86                        Some(server) => {
87                            if server.is_null() {
88                                return Err(WhoIsError::MapError(
89                                    "`ip` in the `_` object in the server list is null.",
90                                ));
91                            }
92
93                            WhoIsServerValue::from_value(server)?
94                        },
95                        None => {
96                            return Err(WhoIsError::MapError(
97                                "Cannot find `ip` in the `_` object in the server list.",
98                            ));
99                        },
100                    }
101                } else {
102                    return Err(WhoIsError::MapError("`_` in the server list is not an object."));
103                }
104            },
105            None => return Err(WhoIsError::MapError("Cannot find `_` in the server list.")),
106        };
107
108        let mut new_map: HashMap<String, WhoIsServerValue> = HashMap::with_capacity(map.len());
109
110        for (k, v) in map {
111            if !v.is_null() {
112                let server_value = WhoIsServerValue::from_value(&v)?;
113                new_map.insert(k, server_value);
114            }
115        }
116
117        Ok(WhoIs {
118            map: new_map,
119            ip,
120        })
121    }
122}
123
124impl WhoIs {
125    pub fn can_find_server_for_tld<T: AsRef<str>, D: AsRef<str>>(
126        &mut self,
127        tld: T,
128        dns_server: D,
129    ) -> bool {
130        let mut tld = tld.as_ref();
131        let dns_server = dns_server.as_ref();
132
133        let address = dns_server.parse().unwrap();
134        let conn = UdpClientConnection::new(address).unwrap();
135        let client = SyncClient::new(conn);
136
137        loop {
138            if self.map.contains_key(tld) {
139                break;
140            }
141
142            match tld.find('.') {
143                Some(index) => {
144                    tld = &tld[index + 1..];
145                },
146                None => {
147                    tld = "";
148                },
149            }
150
151            if tld.is_empty() {
152                break;
153            }
154
155            let name = Name::from_str(&format!("_nicname._tcp.{}.", tld)).unwrap();
156            let response: DnsResponse = client.query(&name, DNSClass::IN, RecordType::SRV).unwrap();
157            let answers: &[Record] = response.answers();
158
159            for record in answers {
160                if let Some(RData::SRV(record)) = record.data() {
161                    let target = record.target().to_string();
162                    let new_server =
163                        match WhoIsServerValue::from_string(&target[..target.len() - 1]) {
164                            Ok(new_server) => new_server,
165                            Err(_error) => continue,
166                        };
167
168                    self.map.insert(tld.to_string(), new_server);
169
170                    return true;
171                }
172            }
173        }
174
175        false
176    }
177
178    fn get_server_by_tld(&self, mut tld: &str) -> Option<&WhoIsServerValue> {
179        let mut server;
180
181        loop {
182            server = self.map.get(tld);
183
184            if server.is_some() {
185                break;
186            }
187
188            if tld.is_empty() {
189                break;
190            }
191
192            match tld.find('.') {
193                Some(index) => {
194                    tld = &tld[index + 1..];
195                },
196                None => {
197                    tld = "";
198                },
199            }
200        }
201
202        server
203    }
204
205    fn lookup_once(
206        server: &WhoIsServerValue,
207        text: &str,
208        timeout: Option<Duration>,
209    ) -> Result<(String, String), WhoIsError> {
210        let addr = server.host.to_addr_string(DEFAULT_WHOIS_HOST_PORT);
211
212        let mut client = if let Some(timeout) = timeout {
213            let socket_addrs: Vec<SocketAddr> = addr.to_socket_addrs()?.collect();
214
215            let mut client = None;
216
217            for socket_addr in socket_addrs.iter().take(socket_addrs.len() - 1) {
218                if let Ok(c) = TcpStream::connect_timeout(socket_addr, timeout) {
219                    client = Some(c);
220                    break;
221                }
222            }
223
224            let client = if let Some(client) = client {
225                client
226            } else {
227                let socket_addr = &socket_addrs[socket_addrs.len() - 1];
228                TcpStream::connect_timeout(socket_addr, timeout)?
229            };
230
231            client.set_read_timeout(Some(timeout))?;
232            client.set_write_timeout(Some(timeout))?;
233            client
234        } else {
235            TcpStream::connect(&addr)?
236        };
237
238        if let Some(query) = &server.query {
239            client.write_all(query.replace("$addr", text).as_bytes())?;
240        } else {
241            client.write_all(DEFAULT_WHOIS_HOST_QUERY.replace("$addr", text).as_bytes())?;
242        }
243
244        client.flush()?;
245
246        let mut query_result = String::new();
247
248        client.read_to_string(&mut query_result)?;
249
250        Ok((addr, query_result))
251    }
252
253    fn lookup_inner(
254        server: &WhoIsServerValue,
255        text: &str,
256        timeout: Option<Duration>,
257        mut follow: u16,
258    ) -> Result<String, WhoIsError> {
259        let mut query_result = Self::lookup_once(server, text, timeout)?;
260
261        while follow > 0 {
262            if let Some(c) = RE_SERVER.captures(&query_result.1) {
263                if let Some(h) = c.get(3) {
264                    let h = h.as_str();
265                    if h.ne(&query_result.0) {
266                        if let Ok(server) = WhoIsServerValue::from_string(h) {
267                            query_result = Self::lookup_once(&server, text, timeout)?;
268
269                            follow -= 1;
270
271                            continue;
272                        }
273                    }
274                }
275            }
276
277            break;
278        }
279
280        Ok(query_result.1)
281    }
282
283    /// Lookup a domain or an IP.
284    pub fn lookup(&self, options: WhoIsLookupOptions) -> Result<String, WhoIsError> {
285        match &options.target.0 {
286            Host::IPv4(_) | Host::IPv6(_) => {
287                let server = match &options.server {
288                    Some(server) => server,
289                    None => &self.ip,
290                };
291
292                Self::lookup_inner(
293                    server,
294                    options.target.to_uri_authority_string().as_ref(),
295                    options.timeout,
296                    options.follow,
297                )
298            },
299            Host::Domain(domain) => {
300                let server = match &options.server {
301                    Some(server) => server,
302                    None => match self.get_server_by_tld(domain.as_str()) {
303                        Some(server) => server,
304                        None => {
305                            return Err(WhoIsError::MapError(
306                                "No whois server is known for this kind of object.",
307                            ));
308                        },
309                    },
310                };
311
312                // punycode check is not necessary because the domain has been ascii-encoded
313
314                Self::lookup_inner(server, domain, options.timeout, options.follow)
315            },
316        }
317    }
318}
319
320#[cfg(feature = "tokio")]
321impl WhoIs {
322    async fn lookup_inner_once_async<'a>(
323        server: &WhoIsServerValue,
324        text: &str,
325        timeout: Option<Duration>,
326    ) -> Result<(String, String), WhoIsError> {
327        let addr = server.host.to_addr_string(DEFAULT_WHOIS_HOST_PORT);
328
329        if let Some(timeout) = timeout {
330            let socket_addrs: Vec<SocketAddr> = addr.to_socket_addrs()?.collect();
331
332            let mut client = None;
333
334            for socket_addr in socket_addrs.iter().take(socket_addrs.len() - 1) {
335                if let Ok(c) =
336                    tokio::time::timeout(timeout, tokio::net::TcpStream::connect(&socket_addr))
337                        .await?
338                {
339                    client = Some(c);
340                    break;
341                }
342            }
343
344            let mut client = if let Some(client) = client {
345                client
346            } else {
347                let socket_addr = &socket_addrs[socket_addrs.len() - 1];
348                tokio::time::timeout(timeout, tokio::net::TcpStream::connect(socket_addr)).await??
349            };
350
351            if let Some(query) = &server.query {
352                tokio::time::timeout(
353                    timeout,
354                    client.write_all(query.replace("$addr", text).as_bytes()),
355                )
356                .await??;
357            } else {
358                tokio::time::timeout(
359                    timeout,
360                    client.write_all(DEFAULT_WHOIS_HOST_QUERY.replace("$addr", text).as_bytes()),
361                )
362                .await??;
363            }
364
365            tokio::time::timeout(timeout, client.flush()).await??;
366
367            let mut query_result = String::new();
368
369            tokio::time::timeout(timeout, client.read_to_string(&mut query_result)).await??;
370
371            Ok((addr, query_result))
372        } else {
373            let mut client = tokio::net::TcpStream::connect(&addr).await?;
374
375            if let Some(query) = &server.query {
376                client.write_all(query.replace("$addr", text).as_bytes()).await?;
377            } else {
378                client
379                    .write_all(DEFAULT_WHOIS_HOST_QUERY.replace("$addr", text).as_bytes())
380                    .await?;
381            }
382
383            client.flush().await?;
384
385            let mut query_result = String::new();
386
387            client.read_to_string(&mut query_result).await?;
388
389            Ok((addr, query_result))
390        }
391    }
392
393    async fn lookup_inner_async<'a>(
394        server: &'a WhoIsServerValue,
395        text: &'a str,
396        timeout: Option<Duration>,
397        mut follow: u16,
398    ) -> Result<String, WhoIsError> {
399        let mut query_result = Self::lookup_inner_once_async(server, text, timeout).await?;
400
401        while follow > 0 {
402            if let Some(c) = RE_SERVER.captures(&query_result.1) {
403                if let Some(h) = c.get(3) {
404                    let h = h.as_str();
405                    if h.ne(&query_result.0) {
406                        if let Ok(server) = WhoIsServerValue::from_string(h) {
407                            query_result =
408                                Self::lookup_inner_once_async(&server, text, timeout).await?;
409
410                            follow -= 1;
411
412                            continue;
413                        }
414                    }
415                }
416            }
417
418            break;
419        }
420
421        Ok(query_result.1)
422    }
423
424    /// Lookup a domain or an IP.
425    pub async fn lookup_async(&self, options: WhoIsLookupOptions) -> Result<String, WhoIsError> {
426        match &options.target.0 {
427            Host::IPv4(_) | Host::IPv6(_) => {
428                let server = match &options.server {
429                    Some(server) => server,
430                    None => &self.ip,
431                };
432
433                Self::lookup_inner_async(
434                    server,
435                    options.target.to_uri_authority_string().as_ref(),
436                    options.timeout,
437                    options.follow,
438                )
439                .await
440            },
441            Host::Domain(domain) => {
442                let server = match &options.server {
443                    Some(server) => server,
444                    None => match self.get_server_by_tld(domain.as_str()) {
445                        Some(server) => server,
446                        None => {
447                            return Err(WhoIsError::MapError(
448                                "No whois server is known for this kind of object.",
449                            ));
450                        },
451                    },
452                };
453
454                // punycode check is not necessary because the domain has been ascii-encoded
455
456                Self::lookup_inner_async(server, domain, options.timeout, options.follow).await
457            },
458        }
459    }
460}
whois_rust/who_is.rs

whois_rust/
who_is.rs
