From c2bd6810b7d473c09b11672b9275c7c3f6ac438e Mon Sep 17 00:00:00 2001
From: Whispering Wind <151555003+ReturnFI@users.noreply.github.com>
Date: Fri, 10 Jan 2025 13:07:24 +0330
Subject: [PATCH] Added GeoCountry

---
 core/scripts/hysteria2/update_geo.py | 79 ++++++++++++++++++----------
 1 file changed, 52 insertions(+), 27 deletions(-)

diff --git a/core/scripts/hysteria2/update_geo.py b/core/scripts/hysteria2/update_geo.py
index 3b227c3..455ab7a 100644
--- a/core/scripts/hysteria2/update_geo.py
+++ b/core/scripts/hysteria2/update_geo.py
@@ -1,46 +1,71 @@
 #!/usr/bin/env python3
 import os
 import subprocess
+from enum import Enum
+import sys
+import requests
+
+
+class GeoCountry(Enum):
+    IRAN = {
+        'geosite': 'https://github.com/chocolate4u/Iran-v2ray-rules/releases/latest/download/geosite.dat',
+        'geoip': 'https://github.com/chocolate4u/Iran-v2ray-rules/releases/latest/download/geoip.dat'
+    }
+    CHINA = {
+        'geosite': 'https://github.com/Loyalsoldier/v2ray-rules-dat/releases/latest/download/geosite.dat',
+        'geoip': 'https://github.com/Loyalsoldier/v2ray-rules-dat/releases/latest/download/geoip.dat'
+    }
+    RUSSIA = {
+        'geosite': 'https://github.com/runetfreedom/russia-v2ray-rules-dat/releases/latest/download/geosite.dat',
+        'geoip': 'https://github.com/runetfreedom/russia-v2ray-rules-dat/releases/latest/download/geoip.dat'
+    }
 
 GEOSITE_PATH = "/etc/hysteria/geosite.dat"
 GEOIP_PATH = "/etc/hysteria/geoip.dat"
-GEOSITE_URL = "https://raw.githubusercontent.com/Chocolate4U/Iran-v2ray-rules/release/geosite.dat"
-GEOIP_URL = "https://raw.githubusercontent.com/Chocolate4U/Iran-v2ray-rules/release/geoip.dat"
-
 
 def remove_file(file_path):
-    if os.path.exists(file_path):
-        os.remove(file_path)
-        print(f"Removed existing file: {file_path}")
-
-
-def download_file(url, destination):
     try:
-        subprocess.run(
-            ["wget", "-O", destination, url],
-            stdout=subprocess.DEVNULL,
-            stderr=subprocess.DEVNULL,
-            check=True
-        )
-        print(f"Downloaded {url} to {destination}")
-    except subprocess.CalledProcessError:
-        print(f"Failed to download {url}")
+        if os.path.exists(file_path):
+            os.remove(file_path)
+            print(f"Removed existing file: {file_path}")
+    except Exception as e:
+        print(f"Error removing file {file_path}: {e}")
 
-
-def update_geo_files():
+def download_file(url, destination, chunk_size=8192):
     try:
-        print("Starting geo files update...")
+        # Ensure the destination directory exists
+        destination_dir = os.path.dirname(destination)
+        if destination_dir and not os.path.exists(destination_dir):
+            os.makedirs(destination_dir)  # Create the directory if it doesn't exist
 
+        response = requests.get(url, stream=True)
+        response.raise_for_status()
+
+        with open(destination, "wb") as file:
+            for chunk in response.iter_content(chunk_size=chunk_size):
+                file.write(chunk)
+
+        print(f"File successfully downloaded to: {destination}")
+
+    except requests.exceptions.RequestException as e:
+        print(f"Error: Failed to download the file from '{url}'.\n{e}")
+    except IOError as e:
+        print(f"Error: Failed to save the file to '{destination}'.\n{e}")
+
+def update_geo_files(country='iran'):
+    try:
+        print(f"Starting geo files update for {country.upper()}...")
+        country_enum = GeoCountry[country.upper()]
         remove_file(GEOSITE_PATH)
         remove_file(GEOIP_PATH)
-        download_file(GEOSITE_URL, GEOSITE_PATH)
-        download_file(GEOIP_URL, GEOIP_PATH)
-
+        download_file(country_enum.value['geosite'], GEOSITE_PATH)
+        download_file(country_enum.value['geoip'], GEOIP_PATH)
         print("Geo files update completed successfully.")
-
+    except KeyError:
+        print(f"Invalid country selection. Available options: {', '.join([c.name.lower() for c in GeoCountry])}")
     except Exception as e:
         print(f"An error occurred during the update process: {e}")
 
-
 if __name__ == "__main__":
-    update_geo_files()
+    country = sys.argv[1] if len(sys.argv) > 1 else 'iran'
+    update_geo_files(country)