Addressing feedback-2

Author: abuabraham-ttd

scripts/confidential_compute.py

@@ -3,7 +3,8 @@
 import socket
 from urllib.parse import urlparse
 from abc import ABC, abstractmethod
-from typing import TypedDict, Dict
+from typing import TypedDict
+
 
 class OperatorConfig(TypedDict):
     enclave_memory_mb: int
@@ -13,70 +14,69 @@ class OperatorConfig(TypedDict):
     core_base_url: str
     optout_base_url: str
 
-class ConfidentialCompute(ABC):
 
+class ConfidentialCompute(ABC):
     @abstractmethod
-    def _get_secret(self, secret_identifier):
+    def _get_secret(self, secret_identifier: str) -> OperatorConfig:
         """
-        Gets the secret from secret store
+        Fetches the secret from a secret store.
 
-        Raises: 
-            SecretNotFoundException: Points to public documentation
+        Raises:
+            SecretNotFoundException: If the secret is not found.
         """
         pass
 
-    def validate_operator_key(self, secrets: OperatorConfig):
-        """
-        Validates operator key if following new pattern. Ignores otherwise
-        """
-        api_token = secrets.get('api_token', None)
+    def validate_operator_key(self, secrets: OperatorConfig) -> bool:
+        """ Validates the operator key format and its environment alignment."""
+        api_token = secrets.get("api_token")
+        if not api_token:
+            raise ValueError("API token is missing from the configuration.")
+
         pattern = r"^(UID2|EUID)-.\-(I|P)-\d+-\*$"
-        if bool(re.match(pattern, api_token)):
-            if secrets.get('debug_mode', False) or secrets.get('environment') == 'integ':
-                if api_token.split('-')[2] != 'I':
-                    raise Exception("Operator key does not match the environment")
-            else:
-                if api_token.split('-')[2] != 'P':
-                    raise Exception("Operator key does not match the environment")
+        if re.match(pattern, api_token):
+            env = secrets.get("environment", "").lower()
+            debug_mode = secrets.get("debug_mode", False)
+            expected_env = "I" if debug_mode or env == "integ" else "P"
+            if api_token.split("-")[2] != expected_env:
+                raise ValueError(
+                    f"Operator key does not match the expected environment ({expected_env})."
+                )
         return True
+    
+    @staticmethod
+    def __resolve_hostname(url: str) -> str:
+        """ Resolves the hostname of a URL to an IP address."""
+        hostname = urlparse(url).netloc
+        return socket.gethostbyname(hostname)
 
-    def validate_connectivity(self, config: OperatorConfig):
-        """
-        Validates core/optout is accessible. 
-        """
+    def validate_connectivity(self, config: OperatorConfig) -> None:
+        """ Validates that the core and opt-out URLs are accessible."""
         try:
-            core_ip = socket.gethostbyname(urlparse(config['core_base_url']).netloc)
-            requests.get(config['core_base_url'], timeout=5)
-            optout_ip = socket.gethostbyname(urlparse(config['optout_base_url']).netloc)
-            requests.get(config['optout_base_url'], timeout=5)
-        except (requests.ConnectionError, requests.Timeout) as e :
-            raise Exception("Failed to reach the URL. -- ERROR CODE, enable IPs? {} {}".format(core_ip, optout_ip), e)
+            core_url = config["core_base_url"]
+            optout_url = config["optout_base_url"]
+            core_ip = self.__resolve_hostname(core_url)
+            requests.get(core_url, timeout=5)
+            optout_ip = self.__resolve_hostname(optout_url)
+            requests.get(optout_url, timeout=5)
+
+        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."
+            )
         except Exception as e:
-            raise Exception("Failed to reach the URL. ")
-        """
-            s3 does not have static IP, and the range returned for s3 is huge to validate. 
-            r = requests.get('https://ip-ranges.amazonaws.com/ip-ranges.json')
-            ips = list(map(lambda x: x['ip_prefix'], filter(lambda x: x['region']=='us-east-1' and x['service'] == 'S3', r.json()['prefixes'])))
-        """
-        return
-    
+            raise Exception("Failed to reach the URLs.") from e
+
     @abstractmethod
-    def _setup_auxiliaries(self):
-        """
-        Sets up auxilary processes required to confidential compute
-        """
+    def _setup_auxiliaries(self) -> None:
+        """ Sets up auxiliary processes required for confidential computing. """
         pass
 
     @abstractmethod
-    def _validate_auxiliaries(self):
-        """
-        Validates auxilary services are running
-        """
+    def _validate_auxiliaries(self) -> None:
+        """ Validates auxiliary services are running."""
         pass
 
     @abstractmethod
-    def run_compute(self):
-        """
-        Runs compute.
-        """
+    def run_compute(self) -> None:
+        """ Runs confidential computing."""
         pass