From a43c5e9020a9b7948636e9f3a34b1ff13b789219 Mon Sep 17 00:00:00 2001
From: Vegard Berg <vgberg@gmail.com>
Date: Thu, 8 Jun 2023 21:54:55 +0200
Subject: [PATCH] The servers used to retrieve public IP via DNS are now
 customizable.

---
 config.sample.toml |  5 +++++
 src/config.rs      | 31 +++++++++++++++++++++++++++++++
 src/dns.rs         | 44 +++++++++++++++++++-------------------------
 src/main.rs        |  2 +-
 4 files changed, 56 insertions(+), 26 deletions(-)

diff --git a/config.sample.toml b/config.sample.toml
index f37f7bb..9424d5e 100644
--- a/config.sample.toml
+++ b/config.sample.toml
@@ -1,4 +1,9 @@
 # api_key = "MY API KEY"
+#
+# [dns]
+# ips = ["208.67.222.222", "208.67.220.220", "2620:119:35::35", "2620:119:53::53"]
+# # Note the final period!!
+# domain = "myip.opendns.com."
 # 
 # [[domain]]
 # name = "mydomain.com"
diff --git a/src/config.rs b/src/config.rs
index 9d6a4d5..a5a2561 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -6,9 +6,40 @@ const SAMPLE_CONFIG: &str = include_str!("../config.sample.toml");
 pub struct Configuration {
     pub api_key: String,
 
+    #[serde(default)]
+    pub dns: Dns,
+
     #[serde(rename = "domain")]
     pub domains: Vec<Domain>,
 }
+
+#[derive(Debug,Serialize, Deserialize, PartialEq)]
+pub struct Dns {
+    #[serde(default = "default_ips")]
+    pub ips: Vec<String>,
+    #[serde(default = "default_domain")]
+    pub domain: String,
+}
+
+impl Default for Dns {
+    fn default() -> Self {
+        Self { ips: default_ips(), domain: default_domain() }
+    }
+}
+
+fn default_ips() -> Vec<String> {
+    vec![
+        String::from("208.67.222.222"),
+        String::from("208.67.220.220"),
+        String::from("2620:119:35::35"),
+        String::from("2620:119:53::53"),
+    ]
+}
+
+fn default_domain() -> String {
+    String::from("myip.opendns.com.")
+}
+
 #[derive(Debug, Serialize, Deserialize, std::cmp::PartialEq)]
 #[allow(non_snake_case)]
 pub struct Domain {
diff --git a/src/dns.rs b/src/dns.rs
index bea08d9..cdba35a 100644
--- a/src/dns.rs
+++ b/src/dns.rs
@@ -3,39 +3,33 @@ use trust_dns_resolver::{
     Resolver,
 };
 
-use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
+use std::{net::{IpAddr, Ipv4Addr, Ipv6Addr}, str::FromStr};
 
-pub fn get_ext_ip() -> anyhow::Result<(Option<Ipv4Addr>, Option<Ipv6Addr>)> {
-    let ip_addrs: Vec<IpAddr> = vec![
-        Ipv4Addr::new(208, 67, 222, 222).into(),
-        Ipv4Addr::new(208, 67, 220, 220).into(),
-        Ipv6Addr::new(0x2620, 0x119, 0x35, 0, 0, 0, 0, 0x35).into(),
-        Ipv6Addr::new(0x2620, 0x119, 0x53, 0, 0, 0, 0, 0x53).into(),
-    ];
+pub fn get_ext_ip(config: &crate::config::Configuration) -> anyhow::Result<(Option<Ipv4Addr>, Option<Ipv6Addr>)> {
+
+    let ip_strs = &config.dns.ips;
+    let domain = config.dns.domain.as_str();
+    let ip_addrs: Vec<IpAddr> = ip_strs.into_iter().filter_map(|ip| IpAddr::from_str(ip).ok()).collect();
 
     let nscg = NameServerConfigGroup::from_ips_clear(ip_addrs.as_slice(), 53, true);
     let resolver_cfg = ResolverConfig::from_parts(None, vec![], nscg);
 
     let resolver = Resolver::new(resolver_cfg, ResolverOpts::default())?;
 
-    let ip4: Vec<Ipv4Addr> = match resolver.ipv4_lookup("myip.opendns.com.") {
-        Ok(ips) => ips.into_iter().collect(),
-        Err(_) => vec![],
+    let ip4= match resolver.ipv4_lookup(domain) {
+        Ok(ips) => ips.into_iter()
+            .collect::<Vec<Ipv4Addr>>()
+            .first()
+            .map(|x| x.to_owned()),
+        Err(_) => None,
     };
 
-    let ip6: Vec<Ipv6Addr> = match resolver.ipv6_lookup("myip.opendns.com.") {
-        Ok(ips) => ips.into_iter().collect(),
-        Err(_) => vec![],
+    let ip6 = match resolver.ipv6_lookup(domain) {
+        Ok(ips) => ips.into_iter()
+            .collect::<Vec<Ipv6Addr>>()
+            .first()
+            .map(|x| x.to_owned()),
+        Err(_) => None,
     };
-
-    let ip4_opt = match ip4.first() {
-        Some(addr) => Some(addr.to_owned()),
-        None => None,
-    };
-    let ip6_opt = match ip6.first() {
-        Some(addr) => Some(addr.to_owned()),
-        None => None,
-    };
-
-    Ok((ip4_opt, ip6_opt))
+    Ok((ip4, ip6))
 }
diff --git a/src/main.rs b/src/main.rs
index 10cae9f..c307a9e 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -22,7 +22,7 @@ fn main() -> anyhow::Result<()> {
     }
     let conf: config::Configuration = config::Configuration::read_config(args.config.as_deref())?;
 
-    let (ip4, ip6) = dns::get_ext_ip()?;
+    let (ip4, ip6) = dns::get_ext_ip(&conf)?;
 
     gandi::update_domains(conf, ip4, ip6)?;
     Ok(())