Merge pull request #159 from spencerwooo:geo-dat

feat: add support for geodata updates

Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "mihoro"
 description = "Mihomo CLI client on Linux."
-version = "0.3.0"
+version = "0.4.0"
 edition = "2021"
 readme = "README.md"
 license = "MIT"
@@ -29,4 +29,4 @@ futures-util = "0.3"
 indicatif = "0.17"
 tokio = { version = "1.34", features = ["full"] }
 truncatable = "0.1"
-anyhow = "1.0.77"
+anyhow = "1.0"

src/cmd.rs

@@ -16,6 +16,8 @@ pub enum Commands {
     Setup,
     /// Update mihomo remote config and restart mihomo.service
     Update,
+    /// Update mihomo geodata
+    UpdateGeodata,
     /// Apply mihomo config overrides and restart mihomo.service
     Apply,
     /// Start mihomo.service with systemctl

src/config.rs

@@ -32,6 +32,10 @@ pub struct MihomoConfig {
     external_controller: Option<String>,
     external_ui: Option<String>,
     secret: Option<String>,
+    pub geodata_mode: Option<bool>,
+    pub geo_auto_update: Option<bool>,
+    pub geo_update_interval: Option<u16>,
+    pub geox_url: Option<GeoxUrl>,
 }
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
@@ -58,6 +62,13 @@ pub enum MihomoLogLevel {
     Debug,
 }
 
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct GeoxUrl {
+    pub geoip: String,
+    pub geosite: String,
+    pub mmdb: String,
+}
+
 impl Config {
     pub fn new() -> Config {
         Config {
@@ -66,6 +77,8 @@ impl Config {
             mihomo_binary_path: String::from("~/.local/bin/mihomo"),
             mihomo_config_root: String::from("~/.config/mihomo"),
             user_systemd_root: String::from("~/.config/systemd/user"),
+
+            // https://wiki.metacubex.one/config/general
             mihomo_config: MihomoConfig {
                 port: 7890,
                 socks_port: 7891,
@@ -77,6 +90,20 @@ impl Config {
                 external_controller: Some(String::from("0.0.0.0:9090")),
                 external_ui: Some(String::from("ui")),
                 secret: None,
+                geodata_mode: Some(false),
+                geo_auto_update: Some(true),
+                geo_update_interval: Some(24),
+                geox_url: Some(GeoxUrl {
+                    geoip: String::from(
+                        "https://cdn.jsdelivr.net/gh/MetaCubeX/meta-rules-dat@release/geoip.dat",
+                    ),
+                    geosite: String::from(
+                        "https://cdn.jsdelivr.net/gh/MetaCubeX/meta-rules-dat@release/geosite.dat",
+                    ),
+                    mmdb: String::from(
+                        "https://cdn.jsdelivr.net/gh/MetaCubeX/meta-rules-dat@release/country.mmdb",
+                    ),
+                }),
             },
         }
     }
@@ -169,6 +196,21 @@ pub struct MihomoYamlConfig {
     #[serde(skip_serializing_if = "Option::is_none")]
     secret: Option<String>,
 
+    #[serde(rename = "geodata-mode", skip_serializing_if = "Option::is_none")]
+    geodata_mode: Option<bool>,
+
+    #[serde(rename = "geo-auto-update", skip_serializing_if = "Option::is_none")]
+    geo_auto_update: Option<bool>,
+
+    #[serde(
+        rename = "geo-update-interval",
+        skip_serializing_if = "Option::is_none"
+    )]
+    geo_update_interval: Option<u16>,
+
+    #[serde(rename = "geox-url", skip_serializing_if = "Option::is_none")]
+    geox_url: Option<GeoxUrl>,
+
     #[serde(flatten)]
     extra: HashMap<String, serde_yaml::Value>,
 }
@@ -196,6 +238,10 @@ pub fn apply_mihomo_override(path: &str, override_config: &MihomoConfig) -> Resu
     mihomo_yaml.external_controller = override_config.external_controller.clone();
     mihomo_yaml.external_ui = override_config.external_ui.clone();
     mihomo_yaml.secret = override_config.secret.clone();
+    mihomo_yaml.geodata_mode = override_config.geodata_mode;
+    mihomo_yaml.geo_auto_update = override_config.geo_auto_update;
+    mihomo_yaml.geo_update_interval = override_config.geo_update_interval;
+    mihomo_yaml.geox_url = override_config.geox_url.clone();
 
     // Write to file
     let serialized_mihomo_yaml = serde_yaml::to_string(&mihomo_yaml)?;

src/main.rs

@@ -35,6 +35,7 @@ async fn cli() -> Result<()> {
     match &args.command {
         Some(Commands::Setup) => mihoro.setup(client).await?,
         Some(Commands::Update) => mihoro.update(client).await?,
+        Some(Commands::UpdateGeodata) => mihoro.update_geodata(client).await?,
         Some(Commands::Apply) => mihoro.apply().await?,
         Some(Commands::Uninstall) => mihoro.uninstall()?,
         Some(Commands::Proxy { proxy }) => mihoro.proxy_commands(proxy)?,

src/mihoro.rs

@@ -44,7 +44,7 @@ impl Mihoro {
         });
     }
 
-    pub async fn setup(self, client: Client) -> Result<()> {
+    pub async fn setup(&self, client: Client) -> Result<()> {
         println!(
             "{} Setting up mihomo's binary, config, and systemd service...",
             &self.prefix.cyan()
@@ -85,6 +85,9 @@ impl Mihoro {
         .await?;
         apply_mihomo_override(&self.mihomo_target_config_path, &self.config.mihomo_config)?;
 
+        // Download geodata
+        self.update_geodata(client).await?;
+
         // Create mihomo.service systemd file
         create_mihomo_service(
             &self.mihomo_target_binary_path,
@@ -98,7 +101,7 @@ impl Mihoro {
         Ok(())
     }
 
-    pub async fn update(self, client: Client) -> Result<()> {
+    pub async fn update(&self, client: Client) -> Result<()> {
         // Download remote mihomo config and apply override
         download_file(
             &client,
@@ -118,7 +121,46 @@ impl Mihoro {
         Ok(())
     }
 
-    pub async fn apply(self) -> Result<()> {
+    pub async fn update_geodata(&self, client: Client) -> Result<()> {
+        if let Some(geox_url) = self.config.mihomo_config.geox_url.clone() {
+            // Download geodata files based on `geodata_mode`
+            let geodata_mode = self.config.mihomo_config.geodata_mode.unwrap_or(false);
+            if geodata_mode {
+                download_file(
+                    &client,
+                    &geox_url.geoip,
+                    format!("{}/geoip.dat", &self.mihomo_target_config_root).as_str(),
+                )
+                .await?;
+                download_file(
+                    &client,
+                    &geox_url.geosite,
+                    format!("{}/geosite.dat", &self.mihomo_target_config_root).as_str(),
+                )
+                .await?;
+            } else {
+                download_file(
+                    &client,
+                    &geox_url.mmdb,
+                    format!("{}/country.mmdb", &self.mihomo_target_config_root).as_str(),
+                )
+                .await?;
+            }
+
+            println!("{} Downloaded and updated geodata", self.prefix.green());
+        } else {
+            println!(
+                "{} `geox_url` undefined, refer to {}",
+                self.prefix.yellow(),
+                "'https://wiki.metacubex.one/config/general/#geo_3'"
+                    .bold()
+                    .underline()
+            );
+        }
+        Ok(())
+    }
+
+    pub async fn apply(&self) -> Result<()> {
         // Apply mihomo config override
         apply_mihomo_override(&self.mihomo_target_config_path, &self.config.mihomo_config).map(
             |_| {
@@ -139,7 +181,7 @@ impl Mihoro {
         Ok(())
     }
 
-    pub fn uninstall(self) -> Result<()> {
+    pub fn uninstall(&self) -> Result<()> {
         Systemctl::new().stop("mihomo.service").execute()?;
         Systemctl::new().disable("mihomo.service").execute()?;
 
@@ -165,7 +207,7 @@ impl Mihoro {
         Ok(())
     }
 
-    pub fn proxy_commands(self, proxy: &Option<ProxyCommands>) -> Result<()> {
+    pub fn proxy_commands(&self, proxy: &Option<ProxyCommands>) -> Result<()> {
         match proxy {
             Some(ProxyCommands::Export) => {
                 println!(