update

Author: clarkxuyang

scripts/aws/ec2.py

@@ -48,12 +48,11 @@ def get_meta_url(cls) -> str:
         return f"http://{cls.AWS_METADATA}/latest/dynamic/instance-identity/document"
 
 
-class EC2(ConfidentialCompute):
+class EC2EntryPoint(ConfidentialCompute):
 
     def __init__(self):
         self.configs: AWSConfidentialComputeConfig = {}
 
-
     def __get_aws_token(self) -> str:
         """Fetches a temporary AWS EC2 metadata token."""
         try:
@@ -85,7 +84,7 @@ def __validate_aws_specific_config(self):
                 if min_capacity.get(key) > int(self.configs.get(key, 10**9)):
                     raise ValueError(f"{key} value ({self.configs.get(key, 0)}) needs to be higher than the minimum required ({min_capacity.get(key)}).")
 
-    def _set_secret(self, secret_identifier: str) -> None:
+    def _set_confidential_config(self, secret_identifier: str) -> None:
         """Fetches a secret value from AWS Secrets Manager and adds defaults"""
 
         def add_defaults(configs: Dict[str, any]) ->  None:
@@ -209,7 +208,7 @@ def __run_nitro_enclave(self):
     def run_compute(self) -> None:
         """Main execution flow for confidential compute."""
         secret_manager_key = self.__get_secret_name_from_userdata()
-        self._set_secret(secret_manager_key)
+        self._set_confidential_config(secret_manager_key)
         print(f"Fetched configs from {secret_manager_key}")
         if not self.configs.get("skip_validations"):
             self.validate_configuration()
@@ -245,13 +244,13 @@ def __kill_auxiliaries(self) -> None:
     parser.add_argument("-o", "--operation", choices=["stop", "start"], default="start", help="Operation to perform.")
     args = parser.parse_args()
     try:
-        ec2 = EC2()
+        ec2 = EC2EntryPoint()
         if args.operation == "stop":
             ec2.cleanup()
         else:
             ec2.run_compute()
     except ConfidentialComputeStartupException as e:
         print("Failed starting up Confidential Compute. Please checks the logs for errors and retry \n", e)
     except Exception as e:
-         print("Unexpected failure while starting up Confidential Compute. Please contact UID support team with this log \n ", e)
+        print("Unexpected failure while starting up Confidential Compute. Please contact UID support team with this log \n ", e)
 

scripts/azure-cc/azureEntryPoint.py

@@ -10,10 +10,9 @@
 import logging
 
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-from confidential_compute import ConfidentialCompute, ConfidentialComputeConfig, MissingConfig, ConfidentialComputeStartupException 
-from azure.identity import DefaultAzureCredential, CredentialUnavailableError
+from confidential_compute import ConfidentialCompute, MissingConfig, SecretAccessException, AuxiliariesException, ConfidentialComputeStartupException 
+from azure.identity import DefaultAzureCredential
 from azure.keyvault.secrets import SecretClient
-from azure.core.exceptions import ResourceNotFoundError, HttpResponseError
 
 class AzureEntryPoint(ConfidentialCompute):
 
@@ -29,41 +28,14 @@ def __init__(self):
         super().__init__()
 
     def __check_env_variables(self):
+        # Check essential env variables
         if AzureEntryPoint.kv_name is None:
             raise MissingConfig(self.__class__.__name__, ["VAULT_NAME"])        
         if AzureEntryPoint.secret_name is None:
             raise MissingConfig(self.__class__.__name__, ["OPERATOR_KEY_SECRET_NAME"])        
         if AzureEntryPoint.env_name is None:
             raise MissingConfig(self.__class__.__name__, ["DEPLOYMENT_ENVIRONMENT"])        
-        logging.info("Env variables validation success")
-    
-    def __set_environment(self):
-        self.configs["environment"] = AzureEntryPoint.env_name
-
-    def _set_secret(self, secret_identifier: str = None):
-        try:
-            credential = DefaultAzureCredential()
-            kv_URL = f"https://{AzureEntryPoint.kv_name}.vault.azure.net"
-            secret_client = SecretClient(vault_url=kv_URL, credential=credential)
-            secret = secret_client.get_secret(AzureEntryPoint.secret_name)
-            # print(f"Secret Value: {secret.value}")
-            self.configs["api_token"] = secret.value
-
-        except CredentialUnavailableError as auth_error:
-            logging.error(f"Read operator key, authentication error: {auth_error}")
-            raise
-
-        except ResourceNotFoundError as not_found_error:
-            logging.error(f"Read operator key, secret not found: {AzureEntryPoint.secret_name}. Error: {not_found_error}")
-            raise
-
-        except HttpResponseError as http_error:
-            logging.error(f"Read operator key, HTTP error occurred: {http_error}")
-            raise
-
-        except Exception as e:
-            logging.error(f"Read operator key, an unexpected error occurred: {e}")
-            raise
+        logging.info("Environment variables validation success")
 
     def __create_final_config(self):      
         TARGET_CONFIG = f"/app/conf/{AzureEntryPoint.env_name}-uid2-config.json"
@@ -93,13 +65,37 @@ def __create_final_config(self):
 
         with open(AzureEntryPoint.FINAL_CONFIG, "r") as file:
             logging.info(file.read())
-    
-    def __set_baseurls(self):
+
+    def __set_base_urls(self):
         with open(AzureEntryPoint.FINAL_CONFIG, "r") as file:
             jdata = json.load(file)
             self.configs["core_base_url"] = jdata["core_attest_url"]
             self.configs["optout_base_url"] = jdata["optout_api_uri"]
 
+    def __set_api_token(self):
+        try:
+            credential = DefaultAzureCredential()
+            kv_URL = f"https://{AzureEntryPoint.kv_name}.vault.azure.net"
+            secret_client = SecretClient(vault_url=kv_URL, credential=credential)
+            secret = secret_client.get_secret(AzureEntryPoint.secret_name)
+            # print(f"Secret Value: {secret.value}")
+            self.configs["api_token"] = secret.value
+
+        except Exception as e:
+            errormsg = f"Read operator key, an unexpected error occurred: {e}"
+            logging.error(errormsg)
+            raise SecretAccessException(self.__class__.__name__, errormsg)
+
+    def _set_confidential_config(self, secret_identifier: str = None):
+        self.configs["skip_validations"] = os.getenv("SKIP_VALIDATIONS", "false").lower() == "true"
+        self.configs["debug_mode"] = os.getenv("DEBUG_MODE", "false").lower() == "true"
+        self.configs["environment"] = AzureEntryPoint.env_name
+
+        # set self.configs["api_token"]
+        self.__set_api_token()
+        # set base urls from final config file
+        self.__set_base_urls()
+
     def __run_operator(self):
 
         # Start the operator
@@ -119,46 +115,46 @@ def __run_operator(self):
         logging.info("-- starting java operator application")
         self.run_command(java_command, separate_process=False)
 
-    def __wait_for_sidecar(self):
+    def _setup_auxiliaries(self):
         logging.info("Waiting for sidecar ...")
 
-        url = "http://169.254.169.254/ping"
+        MAX_RETRIES = 15
+        PING_URL = "http://169.254.169.254/ping"
         delay = 1
-        max_retries = 15
 
-        while True:
+        for attempt in range(1, MAX_RETRIES + 1):
             try:
-                response = requests.get(url, timeout=5)
+                response = requests.get(PING_URL, timeout=5)
                 if response.status_code in [200, 204]:
-                    logging.info("Sidecar started")
+                    logging.info("Sidecar started successfully.")
                     return
                 else:
-                    error_msg = f"Unexpected status code: {response.status_code}, response: {response.text}"
-                    raise Exception(error_msg)
+                    logging.warning(
+                        f"Attempt {attempt}: Unexpected status code {response.status_code}. Response: {response.text}"
+                    )
             except Exception as e:
-                if delay > max_retries:
-                    logging.error(f"Sidecar failed to start after {delay} retries with error {e}", exc_info=True)
-                    sys.exit(1)
-                logging.info(f"Sidecar not started. Retrying in {delay} seconds... {e}")
-                time.sleep(delay)
-                delay += 1
+                logging.info(f"Attempt {attempt}: Error during request - {e}")
+
+            if attempt == MAX_RETRIES:
+                logging.error(
+                    f"Sidecar failed to start after {MAX_RETRIES} attempts. Exiting."
+                )
+                raise AuxiliariesException(self.__class__.__name__)
+
+            logging.info(f"Retrying in {delay} seconds... (Attempt {attempt}/{MAX_RETRIES})")
+            time.sleep(delay)
+            delay += 1
 
     def run_compute(self) -> None:
         """Main execution flow for confidential compute."""
         self.__check_env_variables()
-        self._set_secret()
-        self.__set_environment()
         self.__create_final_config()
-        self.__set_baseurls()
+        self._set_confidential_config()
         if not self.configs.get("skip_validations"):
             self.validate_configuration()
-        self.__wait_for_sidecar()
+        self._setup_auxiliaries()
         self.__run_operator()
 
-    def _setup_auxiliaries(self) -> None:
-        """ Sets up auxiliary processes required for confidential computing. """
-        pass
-
     def _validate_auxiliaries(self) -> None:
         """ Validates auxiliary services are running."""
         pass

scripts/confidential_compute.py

@@ -18,9 +18,9 @@ class ConfidentialComputeConfig(TypedDict):
 class ConfidentialComputeStartupException(Exception):
     def __init__(self, error_name, provider, extra_message=None):
         urls = {
-            "EC2": "https://unifiedid.com/docs/guides/operator-guide-aws-marketplace#uid2-operator-error-codes",
-            "Azure": "https://unifiedid.com/docs/guides/operator-guide-azure-enclave#uid2-operator-error-codes",
-            "GCPEntrypoint": "https://unifiedid.com/docs/guides/operator-private-gcp-confidential-space#uid2-operator-error-codes",
+            "EC2EntryPoint": "https://unifiedid.com/docs/guides/operator-guide-aws-marketplace#uid2-operator-error-codes",
+            "AzureEntryPoint": "https://unifiedid.com/docs/guides/operator-guide-azure-enclave#uid2-operator-error-codes",
+            "GCPEntryPoint": "https://unifiedid.com/docs/guides/operator-private-gcp-confidential-space#uid2-operator-error-codes",
         }
         url = urls.get(provider)
         super().__init__(f"{error_name}\n" + (extra_message if extra_message else "") + f"\nVisit {url} for more details")
@@ -48,7 +48,15 @@ def __init__(self, cls):
 class UID2ServicesUnreachable(ConfidentialComputeStartupException):
     def __init__(self, cls, ip=None):
         super().__init__(error_name=f"E06: {self.__class__.__name__}", provider=cls, extra_message=ip)
-    
+
+class SecretAccessException(ConfidentialComputeStartupException):
+    def __init__(self, cls, inner_message):
+        super().__init__(error_name=f"E07: {self.__class__.__name__}", provider=cls, extra_message=inner_message)
+
+class AuxiliariesException(ConfidentialComputeStartupException):
+    def __init__(self, cls, inner_message = None):
+        super().__init__(error_name=f"E07: {self.__class__.__name__}", provider=cls, extra_message=inner_message)
+
 class ConfidentialCompute(ABC):
 
     def __init__(self):
@@ -112,9 +120,9 @@ def validate_connectivity() -> None:
         logging.info("Completed static validation of confidential compute config values")
 
     @abstractmethod
-    def _set_secret(self, secret_identifier: str) -> None:
+    def _set_confidential_config(self, secret_identifier: str) -> None:
         """
-        Fetches the secret from a secret store.
+        Set ConfidentialComputeConfig
         """
         pass
 

scripts/gcp-oidc/gcp.py

@@ -5,27 +5,25 @@
 from typing import Dict
 import sys
 from google.cloud import secretmanager
-from google.auth import default
-from google.auth.exceptions import DefaultCredentialsError
-from google.api_core.exceptions import PermissionDenied, NotFound
 
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-from confidential_compute import ConfidentialCompute, ConfidentialComputeConfig, MissingConfig, ConfigNotFound, MissingInstanceProfile, ConfidentialComputeStartupException
+from confidential_compute import ConfidentialCompute, ConfidentialComputeConfig, MissingConfig, SecretAccessException, ConfidentialComputeStartupException
 
-class GCPEntrypoint(ConfidentialCompute):
+class GCPEntryPoint(ConfidentialCompute):
 
     def __init__(self):
         super().__init__()
 
-    def _get_secret(self, secret_identifier=None) -> ConfidentialComputeConfig:
+    def _set_confidential_config(self, secret_identifier=None) -> None:
+        
         keys_mapping = {
             "core_base_url": "CORE_BASE_URL",
             "optout_base_url": "OPTOUT_BASE_URL",
             "environment": "DEPLOYMENT_ENVIRONMENT",
             "skip_validations": "SKIP_VALIDATIONS",
             "debug_mode": "DEBUG_MODE",
         }
-        config: ConfidentialComputeConfig = {
+        self.config = {
             key: (os.environ[env_var].lower() == "true" if key in ["skip_validations", "debug_mode"] else os.environ[env_var])
             for key, env_var in keys_mapping.items() if env_var in os.environ
         }
@@ -37,12 +35,9 @@ def _get_secret(self, secret_identifier=None) -> ConfidentialComputeConfig:
             secret_version_name = f"{os.getenv("API_TOKEN_SECRET_NAME")}"
             response = client.access_secret_version(name=secret_version_name)
             secret_value = response.payload.data.decode("UTF-8")
-        except (PermissionDenied, DefaultCredentialsError) as e:
-            raise MissingInstanceProfile(self.__class__.__name__, str(e))
-        except NotFound:
-            raise ConfigNotFound(self.__class__.__name__, f"Secret Manager {os.getenv("API_TOKEN_SECRET_NAME")}")
-        config["api_token"] = secret_value
-        return config
+        except Exception as e:
+            raise SecretAccessException(self.__class__.__name__, str(e))
+        self.config["api_token"] = secret_value
 
     def __populate_operator_config(self, destination):
         target_config = f"/app/conf/{self.configs["environment"].lower()}-config.json"
@@ -63,7 +58,7 @@ def _validate_auxiliaries(self) -> None:
         pass
 
     def run_compute(self) -> None:
-        self.configs = self._get_secret('read_from_env_vars')
+        self._set_confidential_config()
         print(f"Fetched configs")
         if not self.configs.get("skip_validations"):
             self.validate_configuration()
@@ -86,7 +81,7 @@ def run_compute(self) -> None:
 
 if __name__ == "__main__":
     try:
-        gcp = GCP()
+        gcp = GCPEntryPoint()
         gcp.run_compute()
     except ConfidentialComputeStartupException as e:
         print("Failed starting up Confidential Compute. Please checks the logs for errors and retry \n", e)