refactor&delete: simplify AWS VM allocation and remove proxy support

desktop_env/providers/aws/manager.py

@@ -1,7 +1,6 @@
 import os
 from filelock import FileLock
 import boto3
-import psutil
 import logging
 import dotenv
 import signal
@@ -163,60 +162,6 @@ def _allocate_vm(region=DEFAULT_REGION, screen_size=(1920, 1080)):
     return instance_id
 
 
-def _allocate_vm_with_proxy(region=DEFAULT_REGION, proxy_config_file=None, screen_size=(1920, 1080)):
-    """Allocate a VM with proxy configuration"""
-    if not PROXY_SUPPORT_AVAILABLE:
-        logger.warning("Proxy support not available, falling back to regular VM allocation")
-        return _allocate_vm(region, screen_size=screen_size)
-    
-    from desktop_env.providers.aws.provider_with_proxy import AWSProviderWithProxy
-    
-    # Initialize proxy pool if needed
-    if proxy_config_file:
-        init_proxy_pool(proxy_config_file)
-    
-    # Get current proxy
-    proxy_pool = get_global_proxy_pool()
-    current_proxy = proxy_pool.get_next_proxy()
-    
-    if current_proxy:
-        logger.info(f"Allocating VM with proxy: {current_proxy.host}:{current_proxy.port}")
-    
-    # Create provider instance
-    provider = AWSProviderWithProxy(region=region, proxy_config_file=proxy_config_file)
-    
-    # Create new instance
-    instance_id = provider.create_instance_with_proxy(
-        image_id=IMAGE_ID_MAP[region],
-        instance_type=INSTANCE_TYPE,
-        security_groups=[os.getenv('AWS_SECURITY_GROUP_ID')],
-        subnet_id=os.getenv('AWS_SUBNET_ID')
-    )
-    
-    try:
-        ec2_client = boto3.client('ec2', region_name=region)
-        instance_details = ec2_client.describe_instances(InstanceIds=[instance_id])
-        instance = instance_details['Reservations'][0]['Instances'][0]
-        public_ip = instance.get('PublicIpAddress', '')
-        if public_ip:
-            vnc_url = f"http://{public_ip}:5910/vnc.html"
-            logger.info("="*80)
-            logger.info(f"🖥️  VNC Web Access URL: {vnc_url}")
-            logger.info(f"📡 Public IP: {public_ip}")
-            logger.info(f"🆔 Instance ID: {instance_id}")
-            if current_proxy:
-                logger.info(f"🌐 Proxy: {current_proxy.host}:{current_proxy.port}")
-            logger.info("="*80)
-            print(f"\n🌐 VNC Web Access URL: {vnc_url}")
-            if current_proxy:
-                print(f"🔄 Current Proxy: {current_proxy.host}:{current_proxy.port}")
-            print(f"📍 Please open the above address in the browser for remote desktop access\n")
-    except Exception as e:
-        logger.warning(f"Failed to get VNC address for proxy instance {instance_id}: {e}")
-    
-    return instance_id
-
-
 class AWSVMManager(VMManager):
     """
     AWS VM Manager for managing virtual machines on AWS.
@@ -224,15 +169,9 @@ class AWSVMManager(VMManager):
     AWS does not need to maintain a registry of VMs, as it can dynamically allocate and deallocate VMs.
     This class supports both regular VM allocation and proxy-enabled VM allocation.
     """
-    def __init__(self, proxy_config_file=None, **kwargs):
-        self.proxy_config_file = proxy_config_file
+    def __init__(self, **kwargs):
         # self.lock = FileLock(".aws_lck", timeout=60)
         self.initialize_registry()
-        
-        # Initialize proxy pool if proxy configuration is provided
-        if proxy_config_file and PROXY_SUPPORT_AVAILABLE:
-            init_proxy_pool(proxy_config_file)
-            logger.info(f"Proxy pool initialized with config: {proxy_config_file}")
 
     def initialize_registry(self, **kwargs):
         pass
@@ -268,10 +207,6 @@ class AWSVMManager(VMManager):
         pass
 
     def get_vm_path(self, region=DEFAULT_REGION, screen_size=(1920, 1080), **kwargs):
-        if self.proxy_config_file:
-            logger.info("Allocating a new VM with proxy configuration in region: {}".format(region))
-            new_vm_path = _allocate_vm_with_proxy(region, self.proxy_config_file, screen_size=screen_size)
-        else:
-            logger.info("Allocating a new VM in region: {}".format(region))
-            new_vm_path = _allocate_vm(region, screen_size=screen_size)
+        logger.info("Allocating a new VM in region: {}".format(region))
+        new_vm_path = _allocate_vm(region, screen_size=screen_size)
         return new_vm_path

desktop_env/providers/aws/provider_with_proxy.py

@@ -1,315 +0,0 @@
-import boto3
-from botocore.exceptions import ClientError
-import base64
-import logging
-import json
-from typing import Optional
-
-from desktop_env.providers.base import Provider
-from desktop_env.providers.aws.proxy_pool import get_global_proxy_pool, init_proxy_pool, ProxyInfo
-
-logger = logging.getLogger("desktopenv.providers.aws.AWSProviderWithProxy")
-logger.setLevel(logging.INFO)
-
-WAIT_DELAY = 15
-MAX_ATTEMPTS = 10
-
-
-class AWSProviderWithProxy(Provider):
-    
-    def __init__(self, region: str = None, proxy_config_file: str = None):
-        super().__init__(region)
-        self.current_proxy: Optional[ProxyInfo] = None
-        
-        # Initialize proxy pool
-        if proxy_config_file:
-            init_proxy_pool(proxy_config_file)
-            logger.info(f"Initialized proxy pool from {proxy_config_file}")
-        
-        # Get next available proxy
-        self._rotate_proxy()
-
-    def _rotate_proxy(self):
-        """Rotate to next available proxy"""
-        proxy_pool = get_global_proxy_pool()
-        self.current_proxy = proxy_pool.get_next_proxy()
-        
-        if self.current_proxy:
-            logger.info(f"Switched to proxy: {self.current_proxy.host}:{self.current_proxy.port}")
-        else:
-            logger.warning("No proxy available, using direct connection")
-
-    def _generate_proxy_user_data(self) -> str:
-        """Generate user data script with proxy configuration"""
-        if not self.current_proxy:
-            return ""
-        
-        proxy_url = self._format_proxy_url(self.current_proxy)
-        
-        user_data_script = f"""#!/bin/bash
-# Configure system proxy
-echo 'export http_proxy={proxy_url}' >> /etc/environment
-echo 'export https_proxy={proxy_url}' >> /etc/environment
-echo 'export HTTP_PROXY={proxy_url}' >> /etc/environment  
-echo 'export HTTPS_PROXY={proxy_url}' >> /etc/environment
-
-# Configure apt proxy
-cat > /etc/apt/apt.conf.d/95proxy << EOF
-Acquire::http::Proxy "{proxy_url}";
-Acquire::https::Proxy "{proxy_url}";
-EOF
-
-# Configure chrome/chromium proxy
-mkdir -p /etc/opt/chrome/policies/managed
-cat > /etc/opt/chrome/policies/managed/proxy.json << EOF
-{{
-    "ProxyMode": "fixed_servers",
-    "ProxyServer": "{self.current_proxy.host}:{self.current_proxy.port}"
-}}
-EOF
-
-# Configure chromium proxy (Ubuntu default)
-mkdir -p /etc/chromium/policies/managed
-cat > /etc/chromium/policies/managed/proxy.json << EOF
-{{
-    "ProxyMode": "fixed_servers", 
-    "ProxyServer": "{self.current_proxy.host}:{self.current_proxy.port}"
-}}
-EOF
-
-# Configure firefox proxy - support multiple possible paths
-for firefox_dir in /etc/firefox/policies /usr/lib/firefox/distribution/policies /etc/firefox-esr/policies; do
-    if [ -d "$(dirname "$firefox_dir")" ]; then
-        mkdir -p "$firefox_dir"
-        cat > "$firefox_dir/policies.json" << EOF
-{{
-    "policies": {{
-        "Proxy": {{
-            "Mode": "manual",
-            "HTTPProxy": "{self.current_proxy.host}:{self.current_proxy.port}",
-            "HTTPSProxy": "{self.current_proxy.host}:{self.current_proxy.port}",
-            "UseHTTPProxyForAllProtocols": true
-        }}
-    }}
-}}
-EOF
-        break
-    fi
-done
-
-# Reload environment variables
-source /etc/environment
-
-# Log proxy configuration
-echo "$(date): Configured proxy {self.current_proxy.host}:{self.current_proxy.port}" >> /var/log/proxy-setup.log
-"""
-        
-        return base64.b64encode(user_data_script.encode()).decode()
-
-    def _format_proxy_url(self, proxy: ProxyInfo) -> str:
-        """Format proxy URL"""
-        if proxy.username and proxy.password:
-            return f"{proxy.protocol}://{proxy.username}:{proxy.password}@{proxy.host}:{proxy.port}"
-        else:
-            return f"{proxy.protocol}://{proxy.host}:{proxy.port}"
-
-    def start_emulator(self, path_to_vm: str, headless: bool, *args, **kwargs):
-        logger.info("Starting AWS VM with proxy configuration...")
-        ec2_client = boto3.client('ec2', region_name=self.region)
-
-        try:
-            # If instance already exists, start it directly
-            ec2_client.start_instances(InstanceIds=[path_to_vm])
-            logger.info(f"Instance {path_to_vm} is starting...")
-
-            # Wait for the instance to be in the 'running' state
-            waiter = ec2_client.get_waiter('instance_running')
-            waiter.wait(InstanceIds=[path_to_vm], WaiterConfig={'Delay': WAIT_DELAY, 'MaxAttempts': MAX_ATTEMPTS})
-            logger.info(f"Instance {path_to_vm} is now running.")
-
-        except ClientError as e:
-            logger.error(f"Failed to start the AWS VM {path_to_vm}: {str(e)}")
-            raise
-
-    def create_instance_with_proxy(self, image_id: str, instance_type: str, 
-                                 security_groups: list, subnet_id: str) -> str:
-        """Create new instance with proxy configuration"""
-        ec2_client = boto3.client('ec2', region_name=self.region)
-        
-        user_data = self._generate_proxy_user_data()
-        
-        run_instances_params = {
-            "MaxCount": 1,
-            "MinCount": 1,
-            "ImageId": image_id,
-            "InstanceType": instance_type,
-            "EbsOptimized": True,
-            "NetworkInterfaces": [
-                {
-                    "SubnetId": subnet_id,
-                    "AssociatePublicIpAddress": True,
-                    "DeviceIndex": 0,
-                    "Groups": security_groups
-                }
-            ]
-        }
-        
-        if user_data:
-            run_instances_params["UserData"] = user_data
-            
-        try:
-            response = ec2_client.run_instances(**run_instances_params)
-            instance_id = response['Instances'][0]['InstanceId']
-            
-            logger.info(f"Created new instance {instance_id} with proxy configuration")
-            
-            logger.info(f"Waiting for instance {instance_id} to be running...")
-            ec2_client.get_waiter('instance_running').wait(InstanceIds=[instance_id])
-            logger.info(f"Instance {instance_id} is ready.")
-
-            try:
-                instance_details = ec2_client.describe_instances(InstanceIds=[instance_id])
-                instance = instance_details['Reservations'][0]['Instances'][0]
-                public_ip = instance.get('PublicIpAddress', '')
-                if public_ip:
-                    vnc_url = f"http://{public_ip}:5910/vnc.html"
-                    logger.info("="*80)
-                    logger.info(f"🖥️  VNC Web Access URL: {vnc_url}")
-                    logger.info(f"📡 Public IP: {public_ip}")
-                    logger.info(f"🆔 Instance ID: {instance_id}")
-                    if self.current_proxy:
-                        logger.info(f"🌐 Proxy: {self.current_proxy.host}:{self.current_proxy.port}")
-                    logger.info("="*80)
-                    print(f"\n🌐 VNC Web Access URL: {vnc_url}")
-                    if self.current_proxy:
-                        print(f"🔄 Current Proxy: {self.current_proxy.host}:{self.current_proxy.port}")
-                    print(f"📍 Please open the above address in the browser for remote desktop access\n")
-            except Exception as e:
-                logger.warning(f"Failed to get VNC address for instance {instance_id}: {e}")
-            
-            return instance_id
-            
-        except ClientError as e:
-            logger.error(f"Failed to create instance with proxy: {str(e)}")
-            if self.current_proxy:
-                proxy_pool = get_global_proxy_pool()
-                proxy_pool.mark_proxy_failed(self.current_proxy)
-                self._rotate_proxy()
-            raise
-
-    def get_ip_address(self, path_to_vm: str) -> str:
-        logger.info("Getting AWS VM IP address...")
-        ec2_client = boto3.client('ec2', region_name=self.region)
-
-        try:
-            response = ec2_client.describe_instances(InstanceIds=[path_to_vm])
-            for reservation in response['Reservations']:
-                for instance in reservation['Instances']:
-                    private_ip_address = instance.get('PrivateIpAddress', '')
-                    public_ip_address = instance.get('PublicIpAddress', '')
-
-                    if public_ip_address:
-                        vnc_url = f"http://{public_ip_address}:5910/vnc.html"
-                        logger.info("="*80)
-                        logger.info(f"🖥️  VNC Web Access URL: {vnc_url}")
-                        logger.info(f"📡 Public IP: {public_ip_address}")
-                        logger.info(f"🏠 Private IP: {private_ip_address}")
-                        if self.current_proxy:
-                            logger.info(f"🌐 Proxy: {self.current_proxy.host}:{self.current_proxy.port}")
-                        logger.info("="*80)
-                        print(f"\n🌐 VNC Web Access URL: {vnc_url}")
-                        if self.current_proxy:
-                            print(f"🔄 Current Proxy: {self.current_proxy.host}:{self.current_proxy.port}")
-                        print(f"📍 Please open the above address in the browser for remote desktop access\n")
-                    else:
-                        logger.warning("No public IP address available for VNC access")
-                    
-                    return private_ip_address
-            return ''
-        except ClientError as e:
-            logger.error(f"Failed to retrieve IP address for the instance {path_to_vm}: {str(e)}")
-            raise
-
-    def save_state(self, path_to_vm: str, snapshot_name: str):
-        logger.info("Saving AWS VM state...")
-        ec2_client = boto3.client('ec2', region_name=self.region)
-
-        try:
-            image_response = ec2_client.create_image(InstanceId=path_to_vm, Name=snapshot_name)
-            image_id = image_response['ImageId']
-            logger.info(f"AMI {image_id} created successfully from instance {path_to_vm}.")
-            return image_id
-        except ClientError as e:
-            logger.error(f"Failed to create AMI from the instance {path_to_vm}: {str(e)}")
-            raise
-
-    def revert_to_snapshot(self, path_to_vm: str, snapshot_name: str):
-        logger.info(f"Reverting AWS VM to snapshot: {snapshot_name}...")
-        ec2_client = boto3.client('ec2', region_name=self.region)
-
-        try:
-            # Get original instance details for config.
-            instance_details = ec2_client.describe_instances(InstanceIds=[path_to_vm])
-            instance = instance_details['Reservations'][0]['Instances'][0]
-            security_groups = [sg['GroupId'] for sg in instance['SecurityGroups']]
-            subnet_id = instance['SubnetId']
-            instance_type = instance['InstanceType']
-
-            # Terminate the old instance. This is a non-blocking call.
-            logger.info(f"Initiating termination for old instance {path_to_vm}...")
-            ec2_client.terminate_instances(InstanceIds=[path_to_vm])
-            logger.info(f"Old instance {path_to_vm} termination initiated.")
-
-            # Rotate to a new proxy
-            self._rotate_proxy()
-            
-            # Create a new instance
-            new_instance_id = self.create_instance_with_proxy(
-                snapshot_name, instance_type, security_groups, subnet_id
-            )
-            
-            # Note: VNC address is displayed within create_instance_with_proxy
-            logger.info(f"Successfully launched new instance {new_instance_id} for revert.")
-
-            return new_instance_id
-
-        except ClientError as e:
-            logger.error(f"Failed to revert to snapshot {snapshot_name} for the instance {path_to_vm}: {str(e)}")
-            raise
-
-    def stop_emulator(self, path_to_vm, region=None):
-        logger.info(f"Stopping AWS VM {path_to_vm}...")
-        ec2_client = boto3.client('ec2', region_name=self.region)
-
-        try:
-            ec2_client.stop_instances(InstanceIds=[path_to_vm])
-            waiter = ec2_client.get_waiter('instance_stopped')
-            waiter.wait(InstanceIds=[path_to_vm], WaiterConfig={'Delay': WAIT_DELAY, 'MaxAttempts': MAX_ATTEMPTS})
-            logger.info(f"Instance {path_to_vm} has been stopped.")
-        except ClientError as e:
-            logger.error(f"Failed to stop the AWS VM {path_to_vm}: {str(e)}")
-            raise
-
-    def get_current_proxy_info(self) -> Optional[dict]:
-        """Get current proxy information"""
-        if self.current_proxy:
-            return {
-                'host': self.current_proxy.host,
-                'port': self.current_proxy.port,
-                'protocol': self.current_proxy.protocol,
-                'failed_count': self.current_proxy.failed_count
-            }
-        return None
-
-    def force_rotate_proxy(self):
-        """Force rotate proxy"""
-        logger.info("Force rotating proxy...")
-        if self.current_proxy:
-            proxy_pool = get_global_proxy_pool()
-            proxy_pool.mark_proxy_failed(self.current_proxy)
-        self._rotate_proxy()
-
-    def get_proxy_stats(self) -> dict:
-        """Get proxy pool statistics"""
-        proxy_pool = get_global_proxy_pool()
-        return proxy_pool.get_stats()