diff --git a/.github/actions/build_aws_eif/action.yaml b/.github/actions/build_aws_eif/action.yaml index f17523a44..08e6d6604 100644 --- a/.github/actions/build_aws_eif/action.yaml +++ b/.github/actions/build_aws_eif/action.yaml @@ -96,8 +96,9 @@ runs: cp ${{ steps.buildFolder.outputs.BUILD_FOLDER }}/identity_scope.txt ${ARTIFACTS_OUTPUT_DIR}/ cp ${{ steps.buildFolder.outputs.BUILD_FOLDER }}/version_number.txt ${ARTIFACTS_OUTPUT_DIR}/ - cp ./scripts/aws/start.sh ${ARTIFACTS_OUTPUT_DIR}/ - cp ./scripts/aws/stop.sh ${ARTIFACTS_OUTPUT_DIR}/ + cp ./scripts/aws/ec2.py ${ARTIFACTS_OUTPUT_DIR}/ + cp ./scripts/confidential_compute.py ${ARTIFACTS_OUTPUT_DIR}/ + cp ./scripts/aws/requirements.txt ${ARTIFACTS_OUTPUT_DIR}/ cp ./scripts/aws/proxies.host.yaml ${ARTIFACTS_OUTPUT_DIR}/ cp ./scripts/aws/sockd.conf ${ARTIFACTS_OUTPUT_DIR}/ cp ./scripts/aws/uid2operator.service ${ARTIFACTS_OUTPUT_DIR}/ diff --git a/.github/workflows/publish-aws-nitro-eif.yaml b/.github/workflows/publish-aws-nitro-eif.yaml index 8783f6829..31bd87fb4 100644 --- a/.github/workflows/publish-aws-nitro-eif.yaml +++ b/.github/workflows/publish-aws-nitro-eif.yaml @@ -70,7 +70,7 @@ jobs: steps: - name: Build UID2 AWS EIF id: build_uid2_eif - uses: IABTechLab/uid2-operator/.github/actions/build_aws_eif@main + uses: IABTechLab/uid2-operator/.github/actions/build_aws_eif@abu-UID2-4555-EC2-improvements with: identity_scope: uid2 artifacts_base_output_dir: ${{ env.ARTIFACTS_BASE_OUTPUT_DIR }}/uid2 @@ -106,7 +106,7 @@ jobs: steps: - name: Build EUID AWS EIF id: build_euid_eif - uses: IABTechLab/uid2-operator/.github/actions/build_aws_eif@main + uses: IABTechLab/uid2-operator/.github/actions/build_aws_eif@abu-UID2-4555-EC2-improvements with: identity_scope: euid artifacts_base_output_dir: ${{ env.ARTIFACTS_BASE_OUTPUT_DIR }}/euid diff --git a/pom.xml b/pom.xml index d817ec3dd..ba3801aaf 100644 --- a/pom.xml +++ b/pom.xml @@ -6,7 +6,7 @@ com.uid2 uid2-operator - 5.43.0 + 5.43.3-alpha-100-SNAPSHOT UTF-8 diff --git a/scripts/aws/config-server/requirements.txt b/scripts/aws/config-server/requirements.txt index 57652a258..8cdd5ef92 100644 --- a/scripts/aws/config-server/requirements.txt +++ b/scripts/aws/config-server/requirements.txt @@ -1,3 +1,3 @@ Flask==2.3.2 Werkzeug==3.0.3 -setuptools==70.0.0 +setuptools==70.0.0 \ No newline at end of file diff --git a/scripts/aws/ec2.py b/scripts/aws/ec2.py new file mode 100644 index 000000000..b969eb14e --- /dev/null +++ b/scripts/aws/ec2.py @@ -0,0 +1,251 @@ +#!/usr/bin/env python3 + +import boto3 +import json +import os +import subprocess +import re +import multiprocessing +import requests +import signal +import argparse +from botocore.exceptions import ClientError +from typing import Dict +import sys +import time +import yaml + +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 AuxiliaryConfig: + FLASK_PORT: str = "27015" + LOCALHOST: str = "127.0.0.1" + AWS_METADATA: str = "169.254.169.254" + + @classmethod + def get_socks_url(cls) -> str: + return f"socks5://{cls.LOCALHOST}:3306" + + @classmethod + def get_config_url(cls) -> str: + return f"{cls.LOCALHOST}:{cls.FLASK_PORT}/getConfig" + + @classmethod + def get_user_data_url(cls) -> str: + return f"http://{cls.AWS_METADATA}/latest/user-data" + + @classmethod + def get_token_url(cls) -> str: + return f"http://{cls.AWS_METADATA}/latest/api/token" + + @classmethod + def get_meta_url(cls) -> str: + return f"http://{cls.AWS_METADATA}/latest/dynamic/instance-identity/document" + + +class EC2(ConfidentialCompute): + + def __init__(self): + super().__init__() + + def __get_aws_token(self) -> str: + """Fetches a temporary AWS EC2 metadata token.""" + try: + response = requests.put( + AuxiliaryConfig.get_token_url(), headers={"X-aws-ec2-metadata-token-ttl-seconds": "3600"}, timeout=2 + ) + return response.text + except requests.RequestException as e: + raise RuntimeError(f"Failed to fetch aws token: {e}") + + def __get_current_region(self) -> str: + """Fetches the current AWS region from EC2 instance metadata.""" + token = self.__get_aws_token() + headers = {"X-aws-ec2-metadata-token": token} + try: + response = requests.get(AuxiliaryConfig.get_meta_url(), headers=headers, timeout=2) + response.raise_for_status() + return response.json()["region"] + except requests.RequestException as e: + raise RuntimeError(f"Failed to fetch region: {e}") + + 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() + min_capacity = {"enclave_memory_mb": 11000, "enclave_cpu_count" : 2 } + 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)}).") + if min_capacity.get(key) > int(secret.get(key, 10**9)): + raise ValueError(f"{key} value ({secret.get(key, 0)}) needs to be higher than the minimum required ({min_capacity.get(key)}).") + + def _get_secret(self, secret_identifier: str) -> AWSConfidentialComputeConfig: + """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.""" + default_capacity = self.__get_max_capacity() + configs.setdefault("enclave_memory_mb", default_capacity["enclave_memory_mb"]) + configs.setdefault("enclave_cpu_count", default_capacity["enclave_cpu_count"]) + configs.setdefault("debug_mode", False) + return configs + + region = self.__get_current_region() + print(f"Running in {region}") + try: + client = boto3.client("secretsmanager", region_name=region) + except Exception as e: + raise RuntimeError("Please use IAM instance profile for your instance that has permission to access Secret Manager") + try: + secret = add_defaults(json.loads(client.get_secret_value(SecretId=secret_identifier)["SecretString"])) + self.__validate_aws_specific_config(secret) + return secret + except ClientError as _: + raise SecretNotFoundException(f"{secret_identifier} in {region}") + + @staticmethod + def __get_max_capacity(): + try: + with open("/etc/nitro_enclaves/allocator.yaml", "r") as file: + nitro_config = yaml.safe_load(file) + return {"enclave_memory_mb": nitro_config['memory_mib'], "enclave_cpu_count": nitro_config['cpu_count']} + except Exception as e: + raise RuntimeError("/etc/nitro_enclaves/allocator.yaml does not have CPU, memory allocated") + + def __setup_vsockproxy(self, log_level: int) -> None: + """ + Sets up the vsock proxy service. + """ + thread_count = (multiprocessing.cpu_count() + 1) // 2 + command = [ + "/usr/bin/vsockpx", "-c", "/etc/uid2operator/proxy.yaml", + "--workers", str(thread_count), "--log-level", str(log_level), "--daemon" + ] + self.run_command(command) + + def __run_config_server(self) -> None: + """ + Starts the Flask configuration server. + """ + os.makedirs("/etc/secret/secret-value", exist_ok=True) + config_path = "/etc/secret/secret-value/config" + with open(config_path, 'w') as config_file: + json.dump(self.configs, config_file) + os.chdir("/opt/uid2operator/config-server") + command = ["./bin/flask", "run", "--host", AuxiliaryConfig.LOCALHOST, "--port", AuxiliaryConfig.FLASK_PORT] + self.run_command(command, seperate_process=True) + + def __run_socks_proxy(self) -> None: + """ + Starts the SOCKS proxy service. + """ + command = ["sockd", "-D"] + self.run_command(command) + + def __get_secret_name_from_userdata(self) -> str: + """Extracts the secret name from EC2 user data.""" + token = self.__get_aws_token() + response = requests.get(AuxiliaryConfig.get_user_data_url(), headers={"X-aws-ec2-metadata-token": token}) + user_data = response.text + + with open("/opt/uid2operator/identity_scope.txt") as file: + identity_scope = file.read().strip() + + default_name = f"{identity_scope.lower()}-operator-config-key" + hardcoded_value = f"{identity_scope.upper()}_CONFIG_SECRET_KEY" + match = re.search(rf'^export {hardcoded_value}="(.+?)"$', user_data, re.MULTILINE) + return match.group(1) if match else default_name + + def _setup_auxiliaries(self) -> None: + """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() + self.__run_socks_proxy() + + def _validate_auxiliaries(self) -> None: + """Validates connection to flask server direct and through socks proxy.""" + try: + for attempt in range(10): + try: + response = requests.get(AuxiliaryConfig.get_config_url()) + print("Config server is reachable") + break + except requests.exceptions.ConnectionError as e: + print(f"Connecting to config server, attempt {attempt + 1} failed with ConnectionError: {e}") + time.sleep(1) + else: + raise RuntimeError(f"Config server unreachable") + response.raise_for_status() + except requests.RequestException as e: + raise RuntimeError(f"Failed to get config from config server: {e}") + proxies = {"http": AuxiliaryConfig.get_socks_url(), "https": AuxiliaryConfig.get_socks_url()} + try: + response = requests.get(AuxiliaryConfig.get_config_url(), proxies=proxies) + 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_nitro_enclave(self): + command = [ + "nitro-cli", "run-enclave", + "--eif-path", "/opt/uid2operator/uid2operator.eif", + "--memory", str(self.configs["enclave_memory_mb"]), + "--cpu-count", str(self.configs["enclave_cpu_count"]), + "--enclave-cid", "42", + "--enclave-name", "uid2operator" + ] + if self.configs["debug_mode"]: + print("Running in debug_mode") + command += ["--debug-mode", "--attach-console"] + self.run_command(command) + + 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() + self.__run_nitro_enclave() + + def cleanup(self) -> None: + """Terminates the Nitro Enclave and auxiliary processes.""" + try: + self.run_command(["nitro-cli", "terminate-enclave", "--all"]) + self.__kill_auxiliaries() + except subprocess.SubprocessError as e: + raise (f"Error during cleanup: {e}") + + def __kill_auxiliaries(self) -> None: + """Kills a process by its name.""" + try: + for process_name in ["vsockpx", "sockd", "flask"]: + result = subprocess.run(["pgrep", "-f", process_name], stdout=subprocess.PIPE, text=True, check=False) + if result.stdout.strip(): + for pid in result.stdout.strip().split("\n"): + os.kill(int(pid), signal.SIGKILL) + print(f"Killed process '{process_name}'.") + else: + print(f"No process named '{process_name}' found.") + except Exception as e: + print(f"Error killing process '{process_name}': {e}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Manage EC2-based confidential compute workflows.") + parser.add_argument("-o", "--operation", choices=["stop", "start"], default="start", help="Operation to perform.") + args = parser.parse_args() + ec2 = EC2() + if args.operation == "stop": + ec2.cleanup() + else: + ec2.run_compute() + \ No newline at end of file diff --git a/scripts/aws/requirements.txt b/scripts/aws/requirements.txt new file mode 100644 index 000000000..421faba98 --- /dev/null +++ b/scripts/aws/requirements.txt @@ -0,0 +1,4 @@ +requests[socks]==2.32.3 +boto3==1.35.59 +urllib3==1.26.20 +PyYAML===6.0.2 \ No newline at end of file diff --git a/scripts/aws/start.sh b/scripts/aws/start.sh deleted file mode 100644 index 429826928..000000000 --- a/scripts/aws/start.sh +++ /dev/null @@ -1,124 +0,0 @@ -#!/bin/bash - -echo "$HOSTNAME" > /etc/uid2operator/HOSTNAME -EIF_PATH=${EIF_PATH:-/opt/uid2operator/uid2operator.eif} -IDENTITY_SCOPE=${IDENTITY_SCOPE:-$(cat /opt/uid2operator/identity_scope.txt)} -CID=${CID:-42} -TOKEN=$(curl --request PUT "http://169.254.169.254/latest/api/token" --header "X-aws-ec2-metadata-token-ttl-seconds: 3600") -USER_DATA=$(curl -s http://169.254.169.254/latest/user-data --header "X-aws-ec2-metadata-token: $TOKEN") -AWS_REGION_NAME=$(curl -s http://169.254.169.254/latest/dynamic/instance-identity/document/ --header "X-aws-ec2-metadata-token: $TOKEN" | jq -r '.region') -if [ "$IDENTITY_SCOPE" = 'UID2' ]; then - UID2_CONFIG_SECRET_KEY=$([[ "$(echo "${USER_DATA}" | grep UID2_CONFIG_SECRET_KEY=)" =~ ^export\ UID2_CONFIG_SECRET_KEY=\"(.*)\"$ ]] && echo "${BASH_REMATCH[1]}" || echo "uid2-operator-config-key") -elif [ "$IDENTITY_SCOPE" = 'EUID' ]; then - UID2_CONFIG_SECRET_KEY=$([[ "$(echo "${USER_DATA}" | grep EUID_CONFIG_SECRET_KEY=)" =~ ^export\ EUID_CONFIG_SECRET_KEY=\"(.*)\"$ ]] && echo "${BASH_REMATCH[1]}" || echo "euid-operator-config-key") -else - echo "Unrecognized IDENTITY_SCOPE $IDENTITY_SCOPE" - exit 1 -fi -CORE_BASE_URL=$([[ "$(echo "${USER_DATA}" | grep CORE_BASE_URL=)" =~ ^export\ CORE_BASE_URL=\"(.*)\"$ ]] && echo "${BASH_REMATCH[1]}" || echo "") -OPTOUT_BASE_URL=$([[ "$(echo "${USER_DATA}" | grep OPTOUT_BASE_URL=)" =~ ^export\ OPTOUT_BASE_URL=\"(.*)\"$ ]] && echo "${BASH_REMATCH[1]}" || echo "") - -echo "UID2_CONFIG_SECRET_KEY=${UID2_CONFIG_SECRET_KEY}" -echo "CORE_BASE_URL=${CORE_BASE_URL}" -echo "OPTOUT_BASE_URL=${OPTOUT_BASE_URL}" -echo "AWS_REGION_NAME=${AWS_REGION_NAME}" - -function terminate_old_enclave() { - ENCLAVE_ID=$(nitro-cli describe-enclaves | jq -r ".[0].EnclaveID") - [ "$ENCLAVE_ID" != "null" ] && nitro-cli terminate-enclave --enclave-id ${ENCLAVE_ID} -} - -function config_aws() { - aws configure set default.region $AWS_REGION_NAME -} - -function default_cpu() { - target=$(( $(nproc) * 3 / 4 )) - if [ $target -lt 2 ]; then - target="2" - fi - echo $target -} - -function default_mem() { - target=$(( $(grep MemTotal /proc/meminfo | awk '{print $2}') * 3 / 4000 )) - if [ $target -lt 24576 ]; then - target="24576" - fi - echo $target -} - -function read_allocation() { - USER_CUSTOMIZED=$(aws secretsmanager get-secret-value --secret-id "$UID2_CONFIG_SECRET_KEY" | jq -r '.SecretString' | jq -r '.customize_enclave') - shopt -s nocasematch - if [ "$USER_CUSTOMIZED" = "true" ]; then - echo "Applying user customized CPU/Mem allocation..." - CPU_COUNT=${CPU_COUNT:-$(aws secretsmanager get-secret-value --secret-id "$UID2_CONFIG_SECRET_KEY" | jq -r '.SecretString' | jq -r '.enclave_cpu_count')} - MEMORY_MB=${MEMORY_MB:-$(aws secretsmanager get-secret-value --secret-id "$UID2_CONFIG_SECRET_KEY" | jq -r '.SecretString' | jq -r '.enclave_memory_mb')} - else - echo "Applying default CPU/Mem allocation..." - CPU_COUNT=6 - MEMORY_MB=24576 - fi - shopt -u nocasematch -} - - -function update_allocation() { - ALLOCATOR_YAML=/etc/nitro_enclaves/allocator.yaml - if [ -z "$CPU_COUNT" ] || [ -z "$MEMORY_MB" ]; then - echo 'No CPU_COUNT or MEMORY_MB set, cannot start enclave' - exit 1 - fi - echo "updating allocator: CPU_COUNT=$CPU_COUNT, MEMORY_MB=$MEMORY_MB..." - systemctl stop nitro-enclaves-allocator.service - sed -r "s/^(\s*memory_mib\s*:\s*).*/\1$MEMORY_MB/" -i $ALLOCATOR_YAML - sed -r "s/^(\s*cpu_count\s*:\s*).*/\1$CPU_COUNT/" -i $ALLOCATOR_YAML - systemctl start nitro-enclaves-allocator.service && systemctl enable nitro-enclaves-allocator.service - echo "nitro-enclaves-allocator restarted" -} - -function setup_vsockproxy() { - VSOCK_PROXY=${VSOCK_PROXY:-/usr/bin/vsockpx} - VSOCK_CONFIG=${VSOCK_CONFIG:-/etc/uid2operator/proxy.yaml} - VSOCK_THREADS=${VSOCK_THREADS:-$(( ( $(nproc) + 1 ) / 2 )) } - VSOCK_LOG_LEVEL=${VSOCK_LOG_LEVEL:-3} - echo "starting vsock proxy at $VSOCK_PROXY with $VSOCK_THREADS worker threads..." - $VSOCK_PROXY -c $VSOCK_CONFIG --workers $VSOCK_THREADS --log-level $VSOCK_LOG_LEVEL --daemon - echo "vsock proxy now running in background." -} - -function setup_dante() { - sockd -D -} - -function run_config_server() { - mkdir -p /etc/secret/secret-value - { - set +x; # Disable tracing within this block - 2>/dev/null; - SECRET_JSON=$(aws secretsmanager get-secret-value --secret-id "$UID2_CONFIG_SECRET_KEY" | jq -r '.SecretString') - echo "${SECRET_JSON}" > /etc/secret/secret-value/config; - } - echo $(jq ".core_base_url = \"$CORE_BASE_URL\"" /etc/secret/secret-value/config) > /etc/secret/secret-value/config - echo $(jq ".optout_base_url = \"$OPTOUT_BASE_URL\"" /etc/secret/secret-value/config) > /etc/secret/secret-value/config - echo "run_config_server" - cd /opt/uid2operator/config-server - ./bin/flask run --host 127.0.0.1 --port 27015 & -} - -function run_enclave() { - echo "starting enclave..." - nitro-cli run-enclave --eif-path $EIF_PATH --memory $MEMORY_MB --cpu-count $CPU_COUNT --enclave-cid $CID --enclave-name uid2operator -} - -terminate_old_enclave -config_aws -read_allocation -# update_allocation -setup_vsockproxy -setup_dante -run_config_server -run_enclave - -echo "Done!" diff --git a/scripts/aws/stop.sh b/scripts/aws/stop.sh deleted file mode 100644 index c37bdc729..000000000 --- a/scripts/aws/stop.sh +++ /dev/null @@ -1,31 +0,0 @@ -#!/bin/bash - -function terminate_old_enclave() { - echo "Terminating Enclave..." - ENCLAVE_ID=$(nitro-cli describe-enclaves | jq -r ".[0].EnclaveID") - if [ "$ENCLAVE_ID" != "null" ]; then - nitro-cli terminate-enclave --enclave-id $ENCLAVE_ID - else - echo "no running enclaves to terminate" - fi -} - -function kill_process() { - echo "Shutting down $1..." - pid=$(pidof $1) - if [ -z "$pid" ]; then - echo "process $1 not found" - else - kill -9 $pid - echo "$1 exited" - fi -} - -terminate_old_enclave -kill_process vsockpx -kill_process sockd -# we start aws vsock-proxy via nohup -kill_process vsock-proxy -kill_process nohup - -echo "Done!" diff --git a/scripts/aws/uid2-operator-ami/ansible/playbook.yml b/scripts/aws/uid2-operator-ami/ansible/playbook.yml index 84c6c6f14..e6874be94 100644 --- a/scripts/aws/uid2-operator-ami/ansible/playbook.yml +++ b/scripts/aws/uid2-operator-ami/ansible/playbook.yml @@ -70,27 +70,34 @@ 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/start.sh - dest: /opt/uid2operator/start.sh + src: /tmp/artifacts/ec2.py + dest: /opt/uid2operator/ec2.py remote_src: yes - name: Make starter script executable ansible.builtin.file: - path: /opt/uid2operator/start.sh + path: /opt/uid2operator/ec2.py mode: '0755' - - name: Install stopper script + - name: Copy confidential_compute script ansible.builtin.copy: - src: /tmp/artifacts/stop.sh - dest: /opt/uid2operator/stop.sh + src: /tmp/artifacts/confidential_compute.py + dest: /opt/uid2operator/confidential_compute.py remote_src: yes - - name: Make starter script executable - ansible.builtin.file: - path: /opt/uid2operator/stop.sh - mode: '0755' + - name: Create virtualenv for eif init + ansible.builtin.pip: + virtualenv: /opt/uid2operator/init + requirements: /opt/uid2operator/requirements.txt + virtualenv_command: 'python3 -m venv' - name: Install Operator EIF ansible.builtin.copy: diff --git a/scripts/aws/uid2operator.service b/scripts/aws/uid2operator.service index 1d36b7a91..56559e3c2 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=/opt/uid2operator/init/bin/python /opt/uid2operator/ec2.py +ExecStop=/opt/uid2operator/init/bin/python /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 new file mode 100644 index 000000000..5e153d1d5 --- /dev/null +++ b/scripts/confidential_compute.py @@ -0,0 +1,139 @@ +import requests +import re +import socket +from urllib.parse import urlparse +from abc import ABC, abstractmethod +from typing import TypedDict +import subprocess + +class ConfidentialComputeConfig(TypedDict): + debug_mode: bool + api_token: str + core_base_url: str + optout_base_url: str + environment: str + +class ConfidentialCompute(ABC): + + def __init__(self): + self.configs: ConfidentialComputeConfig = {} + + def validate_configuration(self): + """ Validates the paramters specified through configs/secret manager .""" + + def validate_operator_key(): + """ Validates the operator key format and its environment alignment.""" + 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+-.*$" + if re.match(pattern, operator_key): + env = self.configs.get("environment", "").lower() + debug_mode = self.configs.get("debug_mode", False) + expected_env = "I" if debug_mode or env == "integ" else "P" + if operator_key.split("-")[2] != expected_env: + raise ValueError( + f"Operator key does not match the expected environment ({expected_env})." + ) + 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""" + if environment != "prod" and environment not in self.configs[url_key]: + raise ValueError( + f"{url_key} must match the environment. Ensure the URL includes '{environment}'." + ) + parsed_url = urlparse(self.configs[url_key]) + if parsed_url.scheme != 'https' and parsed_url.path: + 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.""" + try: + core_url = self.configs["core_base_url"] + 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." + ) + except Exception as e: + raise Exception("Failed to reach the URLs.") from e + + required_keys = ["api_token", "environment", "core_base_url", "optout_base_url"] + missing_keys = [key for key in required_keys if key not in self.configs] + if missing_keys: + raise ConfidentialComputeMissingConfigError(missing_keys) + + environment = self.configs["environment"] + + if self.configs.get("debug_mode") and environment == "prod": + raise ValueError("Debug mode cannot be enabled in the production environment.") + + validate_url("core_base_url", environment) + validate_url("optout_base_url", environment) + validate_operator_key() + validate_connectivity() + print("Completed static validation of confidential compute config values") + + + @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: + """ Sets up auxiliary processes required for confidential computing. """ + pass + + @abstractmethod + def _validate_auxiliaries(self) -> None: + """ Validates auxiliary services are running.""" + pass + + @abstractmethod + def run_compute(self) -> None: + """ Runs confidential computing.""" + pass + + @staticmethod + def run_command(command, seperate_process=False): + print(f"Running command: {' '.join(command)}") + try: + if seperate_process: + subprocess.Popen(command, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + else: + subprocess.run(command,check=True) + except Exception as e: + print(f"Failed to run command: {str(e)}") + raise RuntimeError (f"Failed to start {' '.join(command)} ") + +class ConfidentialComputeMissingConfigError(Exception): + """Custom exception to handle missing config keys.""" + def __init__(self, missing_keys): + self.missing_keys = missing_keys + self.message = f"Missing configuration keys: {', '.join(missing_keys)}" + super().__init__(self.message) + +class SecretNotFoundException(Exception): + """Custom exception if secret manager is not found""" + def __init__(self, name): + self.message = f"Secret manager not found - {name}" + super().__init__(self.message)