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 # 初始化代理池 if proxy_config_file: init_proxy_pool(proxy_config_file) logger.info(f"Initialized proxy pool from {proxy_config_file}") # 获取下一个可用代理 self._rotate_proxy() def _rotate_proxy(self): """轮换到下一个可用代理""" 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: """生成包含代理配置的user data脚本""" 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: """格式化代理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: # 如果实例已经存在，直接启动 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: """创建带有代理配置的新实例""" 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]: """获取当前代理信息""" 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): """强制轮换代理""" 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: """获取代理池统计信息""" proxy_pool = get_global_proxy_pool() return proxy_pool.get_stats()