From 493abdeeabc8519d05c75f4638e03e4f64d58aa5 Mon Sep 17 00:00:00 2001
From: adlsdztony <zzl0712@connect.hku.hk>
Date: Sat, 7 Jun 2025 11:24:49 +0000
Subject: [PATCH] feat&refactor: add proxy setup functionality and update
 .gitignore for proxy config file

---
 .gitignore                       |  4 ++-
 desktop_env/controllers/setup.py | 54 ++++++++++++++++++++++++++++++++
 desktop_env/desktop_env.py       |  4 +++
 3 files changed, 61 insertions(+), 1 deletion(-)

diff --git a/.gitignore b/.gitignore
index 9ddcd44..48e928c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -195,4 +195,6 @@ vmware_vm_data
 # result
 **/result*/**/*
 
-.vscode
\ No newline at end of file
+.vscode
+
+dataimpulse_proxy_config.json
\ No newline at end of file
diff --git a/desktop_env/controllers/setup.py b/desktop_env/controllers/setup.py
index daf4fe6..096d8e2 100644
--- a/desktop_env/controllers/setup.py
+++ b/desktop_env/controllers/setup.py
@@ -21,11 +21,20 @@ from requests_toolbelt.multipart.encoder import MultipartEncoder
 
 from desktop_env.controllers.python import PythonController
 from desktop_env.evaluators.metrics.utils import compare_urls
+from desktop_env.providers.aws.proxy_pool import get_global_proxy_pool, init_proxy_pool, ProxyInfo
+
+import dotenv
+# Load environment variables from .env file
+dotenv.load_dotenv()
+
+CLIENT_PASSWORD = os.getenv("CLIENT_PASSWORD", "password")  # Default password for sudo operations
+PROXY_CONFIG_FILE = os.getenv("PROXY_CONFIG_FILE", "dataimpulse_proxy_config.json")  # Default proxy config file
 
 logger = logging.getLogger("desktopenv.setup")
 
 FILE_PATH = os.path.dirname(os.path.abspath(__file__))
 
+init_proxy_pool(PROXY_CONFIG_FILE)  # initialize the global proxy pool
 
 class SetupController:
     def __init__(self, vm_ip: str, server_port: int = 5000, chromium_port: int = 9222, vlc_port: int = 8080, cache_dir: str = "cache"):
@@ -360,6 +369,51 @@ class SetupController:
         except requests.exceptions.RequestException as e:
             logger.error("An error occurred while trying to send the request: %s", e)
 
+    def _proxy_setup(self, client_password: str = CLIENT_PASSWORD):
+        """Setup system-wide proxy configuration using proxy pool
+        
+        Args:
+            client_password (str): Password for sudo operations, defaults to "password"
+        """
+        # Get proxy from global proxy pool
+        proxy_pool = get_global_proxy_pool()
+        current_proxy = proxy_pool.get_next_proxy()
+        
+        if not current_proxy:
+            logger.error("No proxy available from proxy pool")
+            raise Exception("No proxy available from proxy pool")
+        
+        # Format proxy URL
+        proxy_url = proxy_pool._format_proxy_url(current_proxy)
+        logger.info(f"Setting up proxy: {current_proxy.host}:{current_proxy.port}")
+        
+        # Configure system proxy environment variables
+        proxy_commands = [
+            f"echo '{client_password}' | sudo -S bash -c \"echo 'export http_proxy={proxy_url}' >> /etc/environment\"",
+            f"echo '{client_password}' | sudo -S bash -c \"echo 'export https_proxy={proxy_url}' >> /etc/environment\"",
+            f"echo '{client_password}' | sudo -S bash -c \"echo 'export HTTP_PROXY={proxy_url}' >> /etc/environment\"",
+            f"echo '{client_password}' | sudo -S bash -c \"echo 'export HTTPS_PROXY={proxy_url}' >> /etc/environment\"",
+        ]
+        
+        # Execute all proxy configuration commands
+        for cmd in proxy_commands:
+            try:
+                self._execute_setup([cmd], shell=True)
+            except Exception as e:
+                logger.error(f"Failed to execute proxy setup command: {e}")
+                proxy_pool.mark_proxy_failed(current_proxy)
+                raise
+        
+        # Reload environment variables
+        reload_cmd = "source /etc/environment"
+        try:
+            logger.info(f"Proxy setup completed successfully for {current_proxy.host}:{current_proxy.port}")
+            proxy_pool.mark_proxy_success(current_proxy)
+        except Exception as e:
+            logger.error(f"Failed to reload environment variables: {e}")
+            proxy_pool.mark_proxy_failed(current_proxy)
+            raise
+
     # Chrome setup
     def _chrome_open_tabs_setup(self, urls_to_open: List[str]):
         host = self.vm_ip
diff --git a/desktop_env/desktop_env.py b/desktop_env/desktop_env.py
index 2a13afa..43114fe 100644
--- a/desktop_env/desktop_env.py
+++ b/desktop_env/desktop_env.py
@@ -191,6 +191,10 @@ class DesktopEnv(gym.Env):
             logger.info("Setting up environment...")
             self.setup_controller.setup(self.config)
             logger.info("Environment setup complete.")
+        
+            if task_config.get("proxy", False):
+                # If using proxy, set up the proxy configuration
+                self.setup_controller._proxy_setup()
 
         observation = self._get_obs()
         return observation