From 8b4600cb63e8ee262d23c7d7e560733b3d601144 Mon Sep 17 00:00:00 2001
From: adlsdztony <zzl0712@outlook.com>
Date: Wed, 28 May 2025 13:28:29 +0800
Subject: [PATCH] feat&refactor: update AWS configuration guidelines and
 improve environment variable handling

---
 desktop_env/providers/aws/AWS_GUIDELINE.md | 17 +++-----
 desktop_env/providers/aws/manager.py       | 51 +++++++++++-----------
 2 files changed, 30 insertions(+), 38 deletions(-)

diff --git a/desktop_env/providers/aws/AWS_GUIDELINE.md b/desktop_env/providers/aws/AWS_GUIDELINE.md
index d45ed0b..4abb158 100644
--- a/desktop_env/providers/aws/AWS_GUIDELINE.md
+++ b/desktop_env/providers/aws/AWS_GUIDELINE.md
@@ -26,18 +26,11 @@ You need to assign values to several variables crucial for the operation of thes
   - Example: `"osworld_key"`
 - **`NETWORK_INTERFACES`**: Configuration settings for network interfaces, which include subnet IDs, security group IDs, and public IP addressing.
   - Example:
-    ```python
-    NETWORK_INTERFACES = {
-        "us-east-1": [
-            {
-                "SubnetId": "subnet-037edfff66c2eb894",
-                "AssociatePublicIpAddress": True,
-                "DeviceIndex": 0,
-                "Groups": ["sg-0342574803206ee9c"]
-            }
-        ],
-        # Add configurations for other regions
-    }
+    ```bash
+    <!-- in .env file -->
+    AWS_REGION=us-east-1
+    AWS_SUBNET_ID=subnet-xxxx
+    AWS_SECURITY_GROUP_ID=sg-xxxx
     ```
 
 
diff --git a/desktop_env/providers/aws/manager.py b/desktop_env/providers/aws/manager.py
index 1e36054..6e8d6de 100644
--- a/desktop_env/providers/aws/manager.py
+++ b/desktop_env/providers/aws/manager.py
@@ -3,6 +3,17 @@ from filelock import FileLock
 import boto3
 import psutil
 import logging
+import dotenv
+# Load environment variables from .env file
+dotenv.load_dotenv()
+
+# Ensure the AWS region is set in the environment
+if not os.getenv('AWS_REGION'):
+    raise EnvironmentError("AWS_REGION must be set in the environment variables.")
+
+# Ensure the AWS subnet and security group IDs are set in the environment
+if not os.getenv('AWS_SUBNET_ID') or not os.getenv('AWS_SECURITY_GROUP_ID'):
+    raise EnvironmentError("AWS_SUBNET_ID and AWS_SECURITY_GROUP_ID must be set in the environment variables.")
 
 from desktop_env.providers.base import VMManager
 
@@ -21,39 +32,27 @@ IMAGE_ID_MAP = {
 
 INSTANCE_TYPE = "t3.medium"
 
-# sg-0342574803206ee9c   subnet-037edfff66c2eb894
-NETWORK_INTERFACE_MAP = {
-    "us-east-1": [
-        {
-            "SubnetId": "subnet-0a4b0c5b8f6066712",
-            "AssociatePublicIpAddress": True,
-            "DeviceIndex": 0,
-            "Groups": [
-                "sg-08a53433e9b4abde6"
-            ]
-        }
-    ],
-    "ap-east-1": [
-        {
-            "SubnetId": "subnet-011060501be0b589c",
-            "AssociatePublicIpAddress": True,
-            "DeviceIndex": 0,
-            "Groups": [
-                "sg-090470e64df78f6eb"
-            ]
-        }
-    ]
-}
-
-
 def _allocate_vm(region=DEFAULT_REGION):
+    
+    if region not in IMAGE_ID_MAP:
+        raise ValueError(f"Region {region} is not supported. Supported regions are: {list(IMAGE_ID_MAP.keys())}")
+
     run_instances_params = {
         "MaxCount": 1,
         "MinCount": 1,
         "ImageId": IMAGE_ID_MAP[region],
         "InstanceType": INSTANCE_TYPE,
         "EbsOptimized": True,
-        "NetworkInterfaces": NETWORK_INTERFACE_MAP[region]
+        "NetworkInterfaces": [
+            {
+                "SubnetId": os.getenv('AWS_SUBNET_ID'),
+                "AssociatePublicIpAddress": True,
+                "DeviceIndex": 0,
+                "Groups": [
+                    os.getenv('AWS_SECURITY_GROUP_ID')
+                ]
+            }
+        ]
     }
 
     ec2_client = boto3.client('ec2', region_name=region)