From 2542232ca4c87842195c05d6b46d6cbe1a563bf6 Mon Sep 17 00:00:00 2001 From: abuabraham-ttd Date: Tue, 10 Dec 2024 11:42:14 -0800 Subject: [PATCH] Add virtual env and start it in systemd --- scripts/aws/ec2.py | 9 +++-- scripts/aws/requirements.txt | 4 +++ .../uid2-operator-ami/ansible/playbook.yml | 16 +++++++++ scripts/aws/uid2operator.service | 6 ++-- scripts/confidential_compute.py | 36 +++++++++---------- 5 files changed, 45 insertions(+), 26 deletions(-) create mode 100644 scripts/aws/requirements.txt diff --git a/scripts/aws/ec2.py b/scripts/aws/ec2.py index db7d5017b..e88f424a2 100644 --- a/scripts/aws/ec2.py +++ b/scripts/aws/ec2.py @@ -47,14 +47,14 @@ def __get_current_region(self) -> str: raise RuntimeError(f"Failed to fetch region: {e}") def __validate_configs(self, secret): - required_keys = ["operator_key", "environment", "core_base_url", "optout_base_url"] + required_keys = ["api_token", "environment", "core_base_url", "optout_base_url"] missing_keys = [key for key in required_keys if key not in secret] if missing_keys: raise ConfidentialComputeMissingConfigError(missing_keys) 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): + 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) -> ConfidentialComputeConfig: @@ -89,7 +89,6 @@ def __add_defaults(self, configs: Dict[str, any]) -> ConfidentialComputeConfig: def __setup_vsockproxy(self, log_level: int) -> None: """ Sets up the vsock proxy service. - TODO: Evaluate adding vsock logging based on log_level here """ thread_count = (multiprocessing.cpu_count() + 1) // 2 command = [ @@ -137,8 +136,8 @@ def _setup_auxiliaries(self) -> None: self.configs = self._get_secret(self.__get_secret_name_from_userdata()) log_level = 3 if self.configs["debug_mode"] else 1 self.__setup_vsockproxy(log_level) - self.__run_config_server(log_level) - self.__run_socks_proxy(log_level) + self.__run_config_server() + self.__run_socks_proxy() time.sleep(5) #TODO: Change to while loop if required. def _validate_auxiliaries(self) -> None: diff --git a/scripts/aws/requirements.txt b/scripts/aws/requirements.txt new file mode 100644 index 000000000..fa562cb02 --- /dev/null +++ b/scripts/aws/requirements.txt @@ -0,0 +1,4 @@ +requests[socks]==2.32.3 +boto3==1.35.59 +urllib3==2.2.3 +PyYAML===5.4.1 \ No newline at end of file diff --git a/scripts/aws/uid2-operator-ami/ansible/playbook.yml b/scripts/aws/uid2-operator-ami/ansible/playbook.yml index 8fb27c2d5..c62c18eee 100644 --- a/scripts/aws/uid2-operator-ami/ansible/playbook.yml +++ b/scripts/aws/uid2-operator-ami/ansible/playbook.yml @@ -70,6 +70,12 @@ requirements: /opt/uid2operator/config-server/requirements.txt virtualenv_command: 'python3 -m venv' + - name: Install requirements.txt for enclave init + ansible.builtin.copy: + src: /tmp/artifacts/requirements.txt + dest: /opt/uid2operator/requirements.txt + remote_src: yes + - name: Install starter script ansible.builtin.copy: src: /tmp/artifacts/ec2.py @@ -87,6 +93,16 @@ dest: /opt/uid2operator/confidential_compute.py remote_src: yes + - name: Install python3-pip + dnf: + name: python3-pip + state: present + + - name: Install dependencies from requirements.txt + pip: + requirements: /opt/uid2operator/requirements.txt + state: present + - name: Install Operator EIF ansible.builtin.copy: src: /tmp/artifacts/uid2operator.eif diff --git a/scripts/aws/uid2operator.service b/scripts/aws/uid2operator.service index 1d36b7a91..18281e2f7 100644 --- a/scripts/aws/uid2operator.service +++ b/scripts/aws/uid2operator.service @@ -8,8 +8,8 @@ RemainAfterExit=true StandardOutput=journal StandardError=journal SyslogIdentifier=uid2operator -ExecStart=/opt/uid2operator/start.sh -ExecStop=/opt/uid2operator/stop.sh +ExecStart=python3 /opt/uid2operator/ec2.py +ExecStop=python3 /opt/uid2operator/ec2.py -o stop [Install] -WantedBy=multi-user.target \ No newline at end of file +WantedBy=multi-user.target diff --git a/scripts/confidential_compute.py b/scripts/confidential_compute.py index a154e37cb..8900ade00 100644 --- a/scripts/confidential_compute.py +++ b/scripts/confidential_compute.py @@ -10,7 +10,7 @@ class ConfidentialComputeConfig(TypedDict): enclave_memory_mb: int enclave_cpu_count: int debug_mode: bool - operator_key: str + api_token: str core_base_url: str optout_base_url: str environment: str @@ -20,16 +20,6 @@ class ConfidentialCompute(ABC): def __init__(self): self.configs: ConfidentialComputeConfig = {} - @abstractmethod - def _get_secret(self, secret_identifier: str) -> ConfidentialComputeConfig: - """ - Fetches the secret from a secret store. - - Raises: - SecretNotFoundException: If the secret is not found. - """ - pass - def validate_environment(self): def validate_url(url_key, environment): if environment not in self.configs[url_key]: @@ -49,7 +39,7 @@ def validate_url(url_key, environment): def validate_operator_key(self): """ Validates the operator key format and its environment alignment.""" - operator_key = self.configs.get("operator_key") + operator_key = self.configs.get("api_token") if not operator_key: raise ValueError("API token is missing from the configuration.") pattern = r"^(UID2|EUID)-.\-(I|P)-\d+-\*$" @@ -62,12 +52,6 @@ def validate_operator_key(self): 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) -> None: """ Validates that the core and opt-out URLs are accessible.""" @@ -84,6 +68,17 @@ def validate_connectivity(self) -> None: ) except Exception as e: raise Exception("Failed to reach the URLs.") from e + + + @abstractmethod + def _get_secret(self, secret_identifier: str) -> ConfidentialComputeConfig: + """ + Fetches the secret from a secret store. + + Raises: + SecretNotFoundException: If the secret is not found. + """ + pass @abstractmethod def _setup_auxiliaries(self) -> None: @@ -100,6 +95,11 @@ def run_compute(self) -> None: """ Runs confidential computing.""" pass + @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) @staticmethod def run_command(command, seperate_process=False):