From aa2ce3e0f4f84575ae4efd510669b47e1ab07362 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?R=C3=A9mi=20BERTHO?= <remi.bertho@dalan.fr>
Date: Wed, 6 Oct 2021 23:09:39 +0200
Subject: [PATCH] Fixed bug

---
 .gitignore  |  3 ++
 TODO.md     |  4 +--
 src/main.rs | 80 ++++++++++++++++++++++++++++++++---------------------
 3 files changed, 53 insertions(+), 34 deletions(-)

diff --git a/.gitignore b/.gitignore
index 52538cc..2954e9e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -18,3 +18,6 @@ Cargo.lock
 /config.toml
 /old_ip.txt
 /log.txt
+
+# IDE
+.idea
diff --git a/TODO.md b/TODO.md
index bde4e85..a9472ec 100644
--- a/TODO.md
+++ b/TODO.md
@@ -1,4 +1,4 @@
 # TODO
  * [ ] Examples de fichiers de config (config + systemd)
- * [ ] N’envoyer qu’un seul mail par test =>_test_domains renvoie Option<String>
- * [ ] Rustfmt
\ No newline at end of file
+ * [ ] Rustfmt
+ * [ ] Fonction check asychrone
\ No newline at end of file
diff --git a/src/main.rs b/src/main.rs
index f4da6bf..b4c4618 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -6,9 +6,9 @@ extern crate log;
 extern crate native_tls;
 extern crate serde;
 extern crate simplelog;
+extern crate tokio;
 extern crate toml;
 extern crate trust_dns_resolver;
-extern crate tokio;
 
 use self::native_tls::{Protocol, TlsConnector};
 use anyhow::{bail, Result};
@@ -17,14 +17,15 @@ use lettre::smtp::authentication::Credentials;
 use lettre::{ClientSecurity, ClientTlsParameters, SmtpClient, Transport};
 use lettre_email::EmailBuilder;
 use log::{error, info};
+use query_external_ip::Consensus;
 use serde::Deserialize;
 use simplelog::{ConfigBuilder, LevelFilter, WriteLogger};
 use std::fs::{self, OpenOptions};
 use std::net::{Ipv4Addr, Ipv6Addr};
+use std::string::String;
 use std::{thread, time::Duration};
-use trust_dns_resolver::{config::*, Name, Resolver};
-use query_external_ip::Consensus;
 use tokio::runtime::{Builder, Runtime};
+use trust_dns_resolver::{config::*, Name, Resolver};
 
 fn main() -> Result<()> {
     // Init log
@@ -82,7 +83,7 @@ fn main() -> Result<()> {
 struct Tester {
     config: Config,
     dns_resolver: Resolver,
-    async_runtime: Runtime
+    async_runtime: Runtime,
 }
 
 impl Tester {
@@ -91,11 +92,8 @@ impl Tester {
         resolver_opt.ip_strategy = LookupIpStrategy::Ipv4AndIpv6;
         Ok(Tester {
             config: toml::from_str(&fs::read_to_string("config.toml")?)?,
-            dns_resolver: Resolver::new(ResolverConfig::default(), resolver_opt)?,
-            async_runtime: Builder::new_multi_thread()
-                .enable_all()
-                .build()
-                .unwrap()
+            dns_resolver: Resolver::from_system_conf()?,
+            async_runtime: Builder::new_multi_thread().enable_all().build().unwrap(),
         })
     }
 
@@ -103,7 +101,10 @@ impl Tester {
         info!("Check");
         match self.get_ip() {
             Ok((ipv4, ipv6)) => {
-                self.test_domains(ipv4, ipv6)?;
+                let msg = self.test_domains(ipv4, ipv6)?;
+                if !msg.is_empty() {
+                    self.send_mail("IP check error".to_owned(), msg)?;
+                }
             }
             Err(err) => {
                 let msg = format!("Cannot get IP ({})", err);
@@ -123,33 +124,48 @@ impl Tester {
         Ok((ipv4, ipv6))
     }
 
-    fn test_domains(&self, ipv4: Ipv4Addr, ipv6: Ipv6Addr) -> Result<()> {
+    fn test_domains(&self, ipv4: Ipv4Addr, ipv6: Ipv6Addr) -> Result<String> {
+        macro_rules! msg {
+        ($full_msg:expr, $($arg:tt)*) => {{
+            let msg = format!($($arg)*);
+            info!("{}", msg);
+            $full_msg.push_str(&msg);
+            $full_msg.push('\n');
+        }}
+}
+
+        let mut full_msg = String::new();
         for domain in self.config.server.domains.iter() {
-            let address = self.dns_resolver.lookup_ip(domain.clone())?;
-            for dns_ip in address.iter() {
-                if dns_ip.is_ipv4() {
-                    if dns_ip != ipv4 {
-                        let msg = format!(
-                            "Wrong IPV4 for {} (DNS: {},  current: {})",
-                            domain, dns_ip, ipv4
-                        );
-                        info!("{}", msg);
-                        self.send_mail("Wrong IP".to_owned(), msg)?;
-                    }
+            if let Ok(dns_ipv4_lookup) = self.dns_resolver.ipv4_lookup(domain.clone()) {
+                let dns_ipv4 = dns_ipv4_lookup.iter().next().unwrap();
+                if *dns_ipv4 != ipv4 {
+                    msg!(
+                        full_msg,
+                        "Wrong IPV4 for {} (DNS: {},  current: {})",
+                        domain,
+                        dns_ipv4,
+                        ipv4
+                    );
                 }
-                if dns_ip.is_ipv6() {
-                    if dns_ip != ipv6 {
-                        let msg = format!(
-                            "Wrong IPV6 for {} (DNS: {},  current: {})",
-                            domain, dns_ip, ipv6
-                        );
-                        info!("{}", msg);
-                        self.send_mail("Wrong IP".to_owned(), msg)?;
-                    }
+            } else {
+                msg!(full_msg, "Cannot get IPv4 address for {}", domain);
+            }
+            if let Ok(dns_ipv6_lookup) = self.dns_resolver.ipv6_lookup(domain.clone()) {
+                let dns_ipv6 = dns_ipv6_lookup.iter().next().unwrap();
+                if *dns_ipv6 != ipv6 {
+                    msg!(
+                        full_msg,
+                        "Wrong IPV6 for {} (DNS: {},  current: {})",
+                        domain,
+                        dns_ipv6,
+                        ipv6
+                    );
                 }
+            } else {
+                msg!(full_msg, "Cannot get IPv6 address for {}", domain);
             }
         }
-        Ok(())
+        Ok(full_msg)
     }
 
     fn send_mail(&self, subject: String, body: String) -> Result<()> {