diff --git a/scripts/aws/ec2.py b/scripts/aws/ec2.py index e5f9ebd19..47d219754 100644 --- a/scripts/aws/ec2.py +++ b/scripts/aws/ec2.py @@ -18,21 +18,20 @@ sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from confidential_compute import ConfidentialCompute, ConfidentialComputeConfig, SecretNotFoundException - class AWSConfidentialComputeConfig(ConfidentialComputeConfig): enclave_memory_mb: int enclave_cpu_count: int - class EC2(ConfidentialCompute): def __init__(self): super().__init__() + self.aws_metadata = "169.254.169.254" def __get_aws_token(self) -> str: """Fetches a temporary AWS EC2 metadata token.""" try: - token_url = "http://169.254.169.254/latest/api/token" + token_url = f"http://{self.aws_metadata}/latest/api/token" response = requests.put( token_url, headers={"X-aws-ec2-metadata-token-ttl-seconds": "3600"}, timeout=2 ) @@ -43,7 +42,7 @@ def __get_aws_token(self) -> str: def __get_current_region(self) -> str: """Fetches the current AWS region from EC2 instance metadata.""" token = self.__get_aws_token() - metadata_url = "http://169.254.169.254/latest/dynamic/instance-identity/document" + metadata_url = f"http://{self.aws_metadata}/latest/dynamic/instance-identity/document" headers = {"X-aws-ec2-metadata-token": token} try: response = requests.get(metadata_url, headers=headers, timeout=2) @@ -55,12 +54,14 @@ def __get_current_region(self) -> str: def __validate_aws_specific_config(self, secret): if "enclave_memory_mb" in secret or "enclave_cpu_count" in secret: max_capacity = self.__get_max_capacity() + for key in ["enclave_memory_mb", "enclave_cpu_count"]: if int(secret.get(key, 0)) > max_capacity.get(key): raise ValueError(f"{key} value ({secret.get(key, 0)}) exceeds the maximum allowed ({max_capacity.get(key)}).") + def _get_secret(self, secret_identifier: str) -> AWSConfidentialComputeConfig: - """Fetches a secret value from AWS Secrets Manager.""" + """Fetches a secret value from AWS Secrets Manager and adds defaults""" def add_defaults(configs: Dict[str, any]) -> AWSConfidentialComputeConfig: """Adds default values to configuration if missing.""" @@ -71,6 +72,7 @@ def add_defaults(configs: Dict[str, any]) -> AWSConfidentialComputeConfig: return configs region = self.__get_current_region() + print(f"Running in {region}") try: client = boto3.client("secretsmanager", region_name=region) except Exception as e: @@ -124,7 +126,7 @@ def __run_socks_proxy(self) -> None: def __get_secret_name_from_userdata(self) -> str: """Extracts the secret name from EC2 user data.""" token = self.__get_aws_token() - user_data_url = "http://169.254.169.254/latest/user-data" + user_data_url = f"http://{self.aws_metadata}/latest/user-data" response = requests.get(user_data_url, headers={"X-aws-ec2-metadata-token": token}) user_data = response.text @@ -137,9 +139,7 @@ def __get_secret_name_from_userdata(self) -> str: return match.group(1) if match else default_name def _setup_auxiliaries(self) -> None: - """Sets up the necessary auxiliary services and configuration.""" - self.configs = self._get_secret(self.__get_secret_name_from_userdata()) - self.validate_configuration() + """Sets up the vsock tunnel, socks proxy and flask server""" log_level = 3 if self.configs["debug_mode"] else 1 self.__setup_vsockproxy(log_level) self.__run_config_server() @@ -147,7 +147,7 @@ def _setup_auxiliaries(self) -> None: time.sleep(5) #TODO: Change to while loop if required. def _validate_auxiliaries(self) -> None: - """Validates auxiliary services.""" + """Validates connection to flask server direct and through socks proxy.""" proxy = "socks5://127.0.0.1:3306" config_url = "http://127.0.0.1:27015/getConfig" try: @@ -161,9 +161,14 @@ def _validate_auxiliaries(self) -> None: response.raise_for_status() except requests.RequestException as e: raise RuntimeError(f"Cannot connect to config server via SOCKS proxy: {e}") + print("Connectivity check to config server passes") def run_compute(self) -> None: """Main execution flow for confidential compute.""" + secret_manager_key = self.__get_secret_name_from_userdata() + self.configs = self._get_secret(secret_manager_key) + print(f"Fetched configs from {secret_manager_key}") + self.validate_configuration() self._setup_auxiliaries() self._validate_auxiliaries() command = [ @@ -175,6 +180,7 @@ def run_compute(self) -> None: "--enclave-name", "uid2operator" ] if self.configs["debug_mode"]: + print("Running in debug_mode") command += ["--debug-mode", "--attach-console"] self.run_command(command) diff --git a/scripts/confidential_compute.py b/scripts/confidential_compute.py index 395c04f49..1769da0af 100644 --- a/scripts/confidential_compute.py +++ b/scripts/confidential_compute.py @@ -35,7 +35,9 @@ def validate_operator_key(): raise ValueError( f"Operator key does not match the expected environment ({expected_env})." ) - return True + print("Validated operator key matches environment") + else: + print("Skipping operator key validation") def validate_url(url_key, environment): """URL should include environment except in prod""" @@ -48,6 +50,8 @@ def validate_url(url_key, environment): raise ValueError( f"{url_key} is invalid. Ensure {self.configs[url_key]} follows HTTPS, and doesn't have any path specified." ) + print(f"Validated {self.configs[url_key]} matches other config parameters") + def validate_connectivity() -> None: """ Validates that the core and opt-out URLs are accessible.""" @@ -56,8 +60,10 @@ def validate_connectivity() -> None: optout_url = self.configs["optout_base_url"] core_ip = socket.gethostbyname(urlparse(core_url).netloc) requests.get(core_url, timeout=5) + print(f"Validated connectivity to {core_url}") optout_ip = socket.gethostbyname(urlparse(optout_url).netloc) requests.get(optout_url, timeout=5) + print(f"Validated connectivity to {optout_url}") except (requests.ConnectionError, requests.Timeout) as e: raise Exception( f"Failed to reach required URLs. Consider enabling {core_ip}, {optout_ip} in the egress firewall." @@ -79,6 +85,7 @@ def validate_connectivity() -> None: validate_url("optout_base_url", environment) validate_operator_key() validate_connectivity() + print("Completed static validation of confidential compute config values") @abstractmethod