275 lines
10 KiB
Python
275 lines
10 KiB
Python
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.")
|
|
|
|
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() |