lybic-guiagents 0.1.0__py3-none-any.whl

desktop_env/providers/aws/provider.py

@@ -0,0 +1,186 @@
+import boto3
+from botocore.exceptions import ClientError
+
+import logging
+
+from desktop_env.providers.base import Provider
+from datetime import datetime
+import time
+
+
+logger = logging.getLogger("desktopenv.providers.aws.AWSProvider")
+logger.setLevel(logging.INFO)
+
+WAIT_DELAY = 15
+MAX_ATTEMPTS = 10
+
+
+class AWSProvider(Provider):
+
+
+    def start_emulator(self, path_to_vm: str, headless: bool, *args, **kwargs):
+        logger.info("Starting AWS VM...")
+        ec2_client = boto3.client('ec2', region_name=self.region)
+
+        try:
+            # Check the current state of the instance
+            response = ec2_client.describe_instances(InstanceIds=[path_to_vm])
+            state = response['Reservations'][0]['Instances'][0]['State']['Name']
+            logger.info(f"Instance {path_to_vm} current state: {state}")
+
+            if state == 'running':
+                # If the instance is already running, skip starting it
+                logger.info(f"Instance {path_to_vm} is already running. Skipping start.")
+                return
+
+            if state == 'stopped':
+                # Start the instance if it's currently stopped
+                ec2_client.start_instances(InstanceIds=[path_to_vm])
+                logger.info(f"Instance {path_to_vm} is starting...")
+
+                # Wait until the instance reaches '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.")
+            else:
+                # For all other states (terminated, pending, etc.), log a warning
+                logger.warning(f"Instance {path_to_vm} is in state '{state}' and cannot be started.")
+
+        except ClientError as e:
+            logger.error(f"Failed to start the AWS VM {path_to_vm}: {str(e)}")
+            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}")
+                        logger.info("="*80)
+                        print(f"\n🌐 VNC Web Access URL: {vnc_url}")
+                        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 ''  # Return an empty string if no IP address is found
+        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 AMI: {snapshot_name}...")
+        ec2_client = boto3.client('ec2', region_name=self.region)
+
+        try:
+            # Step 1: Retrieve the original instance details
+            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']
+            
+            # Step 2: Terminate the old instance
+            ec2_client.terminate_instances(InstanceIds=[path_to_vm])
+            logger.info(f"Old instance {path_to_vm} has been terminated.")
+
+            # Step 3: Launch a new instance from the snapshot(AMI) with performance optimization
+            logger.info(f"Launching a new instance from AMI {snapshot_name}...")
+            
+            run_instances_params = {
+                "MaxCount": 1,
+                "MinCount": 1,
+                "ImageId": snapshot_name,
+                "InstanceType": instance_type,
+                "EbsOptimized": True,
+                "NetworkInterfaces": [
+                    {
+                        "SubnetId": subnet_id,
+                        "AssociatePublicIpAddress": True,
+                        "DeviceIndex": 0,
+                        "Groups": security_groups
+                    }
+                ],
+                "BlockDeviceMappings": [
+                    {
+                        "DeviceName": "/dev/sda1", 
+                        "Ebs": {
+                            # "VolumeInitializationRate": 300
+                            "VolumeSize": 30,  # Size in GB
+                            "VolumeType": "gp3",  # General Purpose SSD
+                            "Throughput": 1000,
+                            "Iops": 4000  # Adjust IOPS as needed
+                        }
+                    }
+                ]
+            }
+            
+            new_instance = ec2_client.run_instances(**run_instances_params)
+            new_instance_id = new_instance['Instances'][0]['InstanceId']
+            logger.info(f"New instance {new_instance_id} launched from AMI {snapshot_name}.")
+            logger.info(f"Waiting for instance {new_instance_id} to be running...")
+            ec2_client.get_waiter('instance_running').wait(InstanceIds=[new_instance_id])
+
+            logger.info(f"Instance {new_instance_id} is ready.")
+            
+            try:
+                instance_details = ec2_client.describe_instances(InstanceIds=[new_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"🖥️  New Instance VNC Web Access URL: {vnc_url}")
+                    logger.info(f"📡 Public IP: {public_ip}")
+                    logger.info(f"🆔 New Instance ID: {new_instance_id}")
+                    logger.info("="*80)
+                    print(f"\n🌐 New Instance VNC Web Access URL: {vnc_url}")
+                    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 new instance {new_instance_id}: {e}")
+
+            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.terminate_instances(InstanceIds=[path_to_vm])
+            logger.info(f"Instance {path_to_vm} has been terminated.")
+        except ClientError as e:
+            logger.error(f"Failed to stop the AWS VM {path_to_vm}: {str(e)}")
+            raise

desktop_env/providers/aws/provider_with_proxy.py

@@ -0,0 +1,315 @@
+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()