feat: implement proxy management for AWS VM provider and enhance task configuration handling
This commit is contained in:
Timothyxxx
261
desktop_env/providers/aws/provider_with_proxy.py
Normal file
261
desktop_env/providers/aws/provider_with_proxy.py
Normal file
@@ -0,0 +1,261 @@
|
||||
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
|
||||
# 配置系统代理
|
||||
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
|
||||
|
||||
# 配置apt代理
|
||||
cat > /etc/apt/apt.conf.d/95proxy << EOF
|
||||
Acquire::http::Proxy "{proxy_url}";
|
||||
Acquire::https::Proxy "{proxy_url}";
|
||||
EOF
|
||||
|
||||
# 配置chrome/chromium代理
|
||||
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
|
||||
|
||||
# 配置firefox代理
|
||||
mkdir -p /etc/firefox/policies
|
||||
cat > /etc/firefox/policies/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
|
||||
|
||||
# 重新加载环境变量
|
||||
source /etc/environment
|
||||
|
||||
# 记录代理配置日志
|
||||
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):
|
||||
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.")
|
||||
|
||||
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', '')
|
||||
return private_ip_address
|
||||
return ''
|
||||
except ClientError as e:
|
||||
logger.error(f"Failed to retrieve private 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:
|
||||
# 获取原实例详情
|
||||
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']
|
||||
|
||||
# 终止旧实例
|
||||
ec2_client.terminate_instances(InstanceIds=[path_to_vm])
|
||||
logger.info(f"Old instance {path_to_vm} has been terminated.")
|
||||
|
||||
# 轮换到新的代理
|
||||
self._rotate_proxy()
|
||||
|
||||
# 创建新实例
|
||||
new_instance_id = self.create_instance_with_proxy(
|
||||
snapshot_name, instance_type, security_groups, subnet_id
|
||||
)
|
||||
|
||||
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()
|
||||
Reference in New Issue
Block a user