From 77b457a14427f3818521f5f5fb1035fe04ff55ac Mon Sep 17 00:00:00 2001 From: Abu Abraham Date: Mon, 18 Nov 2024 11:23:58 -0800 Subject: [PATCH 1/9] P1 --- scripts/aws/Dockerfile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/aws/Dockerfile b/scripts/aws/Dockerfile index e210001c3..6007437b1 100644 --- a/scripts/aws/Dockerfile +++ b/scripts/aws/Dockerfile @@ -45,4 +45,4 @@ COPY ./syslog-ng-client.conf /etc/syslog-ng/syslog-ng.conf RUN chmod +x /app/vsockpx && chmod +x /app/entrypoint.sh -CMD ["/app/entrypoint.sh"] +CMD ["/app/ec2.py"] From 4fc8c157920bf4867cec5764f29d209a3926d31e Mon Sep 17 00:00:00 2001 From: abuabraham-ttd Date: Mon, 18 Nov 2024 11:44:13 -0800 Subject: [PATCH 2/9] Adding pre-init validation, switching to python and interfaces --- scripts/aws/Dockerfile | 2 +- scripts/aws/config-server/requirements.txt | 2 + scripts/aws/ec2.py | 179 ++++++++++++++++++ scripts/aws/start.sh | 124 ------------ scripts/aws/stop.sh | 31 --- .../uid2-operator-ami/ansible/playbook.yml | 17 +- scripts/aws/uid2operator.service | 4 +- scripts/confidential_compute.py | 73 +++++++ 8 files changed, 260 insertions(+), 172 deletions(-) create mode 100755 scripts/aws/ec2.py delete mode 100644 scripts/aws/start.sh delete mode 100644 scripts/aws/stop.sh create mode 100644 scripts/confidential_compute.py diff --git a/scripts/aws/Dockerfile b/scripts/aws/Dockerfile index 6007437b1..e210001c3 100644 --- a/scripts/aws/Dockerfile +++ b/scripts/aws/Dockerfile @@ -45,4 +45,4 @@ COPY ./syslog-ng-client.conf /etc/syslog-ng/syslog-ng.conf RUN chmod +x /app/vsockpx && chmod +x /app/entrypoint.sh -CMD ["/app/ec2.py"] +CMD ["/app/entrypoint.sh"] diff --git a/scripts/aws/config-server/requirements.txt b/scripts/aws/config-server/requirements.txt index 57652a258..bd64bfe35 100644 --- a/scripts/aws/config-server/requirements.txt +++ b/scripts/aws/config-server/requirements.txt @@ -1,3 +1,5 @@ Flask==2.3.2 Werkzeug==3.0.3 setuptools==70.0.0 +requests==2.32.3 +boto3==1.35.59 \ No newline at end of file diff --git a/scripts/aws/ec2.py b/scripts/aws/ec2.py new file mode 100755 index 000000000..27a6cb15f --- /dev/null +++ b/scripts/aws/ec2.py @@ -0,0 +1,179 @@ +#!/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 +import sys +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from confidential_compute import ConfidentialCompute + +class EC2(ConfidentialCompute): + + def __init__(self): + super().__init__() + self.config = {} + + def _get_secret(self, secret_identifier): + client = boto3.client("secretsmanager", region_name=self.__get_current_region()) + try: + secret = client.get_secret_value(SecretId=secret_identifier) + return json.loads(secret["SecretString"]) + except ClientError as e: + raise Exception("Unable to access secret store") + + def __add_defaults(self, configs): + configs.setdefault("enclave_memory_mb", 24576) + configs.setdefault("enclave_cpu_count", 6) + configs.setdefault("debug_mode", False) + return configs + + def __setup_vsockproxy(self, log_level): + thread_count = int((multiprocessing.cpu_count() + 1) // 2) + log_level = log_level + try: + subprocess.Popen(["/usr/bin/vsockpx", "-c", "/etc/uid2operator/proxy.yaml", "--workers", str(thread_count), "--log-level", log_level, "--daemon"]) + print("VSOCK proxy is now running in the background") + except FileNotFoundError: + print("Error: vsockpx not found. Please ensure the path is correct") + except Exception as e: + print("Failed to start VSOCK proxy") + + def __run_config_server(self, log_level): + os.makedirs("/etc/secret/secret-value", exist_ok=True) + with open('/etc/secret/secret-value/config', 'w') as fp: + json.dump(self.configs, fp) + os.chdir("/opt/uid2operator/config-server") + # TODO: Add --log-level to flask. + try: + subprocess.Popen(["./bin/flask", "run", "--host", "127.0.0.1", "--port", "27015"]) + print("Config server is now running in the background.") + except Exception as e: + print(f"Failed to start config server: {e}") + + def __run_socks_proxy(self, log_level): + subprocess.Popen(["sockd", "-d"]) + + def _validate_auxilaries(self): + proxy = "socks5h://127.0.0.1:3305" + url = "http://127.0.0.1:27015/getConfig" + response = requests.get(url) + if response.status_code != 200: + raise Exception("Config server unreachable") + proxies = { + "http": proxy, + "https": proxy, + } + try: + response = requests.get(url, proxies=proxies) + response.raise_for_status() + except Exception as e: + raise Exception(f"Cannot conect to config server through socks5: {e}") + pass + + def __get_aws_token(self): + try: + token_url = "http://169.254.169.254/latest/api/token" + token_response = requests.put(token_url, headers={"X-aws-ec2-metadata-token-ttl-seconds": "3600"}, timeout=2) + return token_response.text + except Exception as e: + return "blank" + + def __get_current_region(self): + token = self.__get_aws_token() + metadata_url = "http://169.254.169.254/latest/dynamic/instance-identity/document" + headers = {"X-aws-ec2-metadata-token": token} + try: + response = requests.get(metadata_url, headers=headers,timeout=2) + if response.status_code == 200: + return response.json().get("region") + else: + print(f"Failed to fetch region, status code: {response.status_code}") + except Exception as e: + raise Exception(f"Region not found, are you running in EC2 environment. {e}") + + def __get_secret_name_from_userdata(self): + token = self.__get_aws_token() + user_data_url = "http://169.254.169.254/latest/user-data" + user_data_response = requests.get(user_data_url, headers={"X-aws-ec2-metadata-token": token}) + user_data = user_data_response.text + identity_scope = open("/opt/uid2operator/identity_scope.txt").read().strip() + default_name = "{}-operator-config-key".format(identity_scope.lower()) + hardcoded_value = "{}_CONFIG_SECRET_KEY".format(identity_scope.upper()) + match = re.search(rf'^export {hardcoded_value}="(.+?)"$', user_data, re.MULTILINE) + return match.group(1) if match else default_name + + def _setup_auxilaries(self): + hostname = os.getenv("HOSTNAME", default=os.uname()[1]) + file_path = "HOSTNAME" + try: + with open(file_path, "w") as file: + file.write(hostname) + print(f"Hostname '{hostname}' written to {file_path}") + except Exception as e: + print(f"An error occurred : {e}") + config = self._get_secret(self.__get_secret_name_from_userdata()) + self.configs = self.__add_defaults(config) + 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) + + def run_compute(self): + self._setup_auxilaries() + self._validate_auxilaries() + command = [ + "nitro-cli", "run-enclave", + "--eif-path", "/opt/uid2operator/uid2operator.eif", + "--memory", self.config['enclave_memory_mb'], + "--cpu-count", self.config['enclave_cpu_count'], + "--enclave-cid", 42, + "--enclave-name", "uid2operator" + ] + if self.config['debug']: + command+=["--debug-mode", "--attach-console"] + subprocess.run(command, check=True) + + def cleanup(self): + describe_output = subprocess.check_output(["nitro-cli", "describe-enclaves"], text=True) + enclaves = json.loads(describe_output) + enclave_id = enclaves[0].get("EnclaveID") if enclaves else None + if enclave_id: + subprocess.run(["nitro-cli", "terminate-enclave", "--enclave-id", enclave_id]) + print(f"Enclave with ID {enclave_id} has been terminated.") + else: + print("No enclave found or EnclaveID is null.") + + def kill_process(self, process_name): + try: + 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"{process_name} exited") + else: + print(f"Process {process_name} not found") + except Exception as e: + print(f"Failed to shut down {process_name}: {e}") + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("-o", "--operation", required=False) + args = parser.parse_args() + ec2 = EC2() + if args.operation and args.operation == "stop": + ec2.cleanup() + [ec2.kill_process(process) for process in ["vsockpx", "sockd", "vsock-proxy", "nohup"]] + else: + ec2.run_compute() 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..3629b34e3 100644 --- a/scripts/aws/uid2-operator-ami/ansible/playbook.yml +++ b/scripts/aws/uid2-operator-ami/ansible/playbook.yml @@ -72,24 +72,13 @@ - 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 - mode: '0755' - - - name: Install stopper script - ansible.builtin.copy: - src: /tmp/artifacts/stop.sh - dest: /opt/uid2operator/stop.sh - remote_src: yes - - - name: Make starter script executable - ansible.builtin.file: - path: /opt/uid2operator/stop.sh + path: /opt/uid2operator/ec2.py mode: '0755' - name: Install Operator EIF diff --git a/scripts/aws/uid2operator.service b/scripts/aws/uid2operator.service index 1d36b7a91..e92f5d401 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/ec2.py +ExecStop=/opt/uid2operator/ec2.py -o stop [Install] WantedBy=multi-user.target \ No newline at end of file diff --git a/scripts/confidential_compute.py b/scripts/confidential_compute.py new file mode 100644 index 000000000..a60220bcd --- /dev/null +++ b/scripts/confidential_compute.py @@ -0,0 +1,73 @@ +import requests +import re +import socket +from urllib.parse import urlparse +from abc import ABC, abstractmethod + +class ConfidentialCompute(ABC): + + @abstractmethod + def _get_secret(self, secret_identifier): + """ + Gets the secret from secret store + + Raises: + SecretNotFoundException: Points to public documentation + """ + pass + + def validate_operator_key(self, secrets): + """ + Validates operator key if following new pattern. Ignores otherwise + """ + api_token = secrets.get('api_token', None) + 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") + return True + + def validate_connectivity(self, config): + """ + Validates core/optout is 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) + 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 + + @abstractmethod + def _setup_auxilaries(self, secrets): + """ + Sets up auxilary processes required to confidential compute + """ + pass + + @abstractmethod + def _validate_auxilaries(self, secrets): + """ + Validates auxilary services are running + """ + pass + + @abstractmethod + def run_compute(self): + """ + Runs compute. + """ + pass \ No newline at end of file From 130340dbe574a0a9dd869a95bcc56059a45ec4db Mon Sep 17 00:00:00 2001 From: abuabraham-ttd Date: Mon, 18 Nov 2024 11:49:04 -0800 Subject: [PATCH 3/9] Adding pre-init validation, switching to python and interfaces --- scripts/aws/config-server/requirements.txt | 3 +- scripts/aws/ec2.py | 78 +++++++++++----------- 2 files changed, 41 insertions(+), 40 deletions(-) diff --git a/scripts/aws/config-server/requirements.txt b/scripts/aws/config-server/requirements.txt index bd64bfe35..5355382c1 100644 --- a/scripts/aws/config-server/requirements.txt +++ b/scripts/aws/config-server/requirements.txt @@ -2,4 +2,5 @@ Flask==2.3.2 Werkzeug==3.0.3 setuptools==70.0.0 requests==2.32.3 -boto3==1.35.59 \ No newline at end of file +boto3==1.35.59 +urllib3==2.2.3 \ No newline at end of file diff --git a/scripts/aws/ec2.py b/scripts/aws/ec2.py index 27a6cb15f..bd4173953 100755 --- a/scripts/aws/ec2.py +++ b/scripts/aws/ec2.py @@ -20,6 +20,27 @@ def __init__(self): super().__init__() self.config = {} + def __get_aws_token(self): + try: + token_url = "http://169.254.169.254/latest/api/token" + token_response = requests.put(token_url, headers={"X-aws-ec2-metadata-token-ttl-seconds": "3600"}, timeout=2) + return token_response.text + except Exception as e: + return "blank" + + def __get_current_region(self): + token = self.__get_aws_token() + metadata_url = "http://169.254.169.254/latest/dynamic/instance-identity/document" + headers = {"X-aws-ec2-metadata-token": token} + try: + response = requests.get(metadata_url, headers=headers,timeout=2) + if response.status_code == 200: + return response.json().get("region") + else: + print(f"Failed to fetch region, status code: {response.status_code}") + except Exception as e: + raise Exception(f"Region not found, are you running in EC2 environment. {e}") + def _get_secret(self, secret_identifier): client = boto3.client("secretsmanager", region_name=self.__get_current_region()) try: @@ -58,45 +79,7 @@ def __run_config_server(self, log_level): print(f"Failed to start config server: {e}") def __run_socks_proxy(self, log_level): - subprocess.Popen(["sockd", "-d"]) - - def _validate_auxilaries(self): - proxy = "socks5h://127.0.0.1:3305" - url = "http://127.0.0.1:27015/getConfig" - response = requests.get(url) - if response.status_code != 200: - raise Exception("Config server unreachable") - proxies = { - "http": proxy, - "https": proxy, - } - try: - response = requests.get(url, proxies=proxies) - response.raise_for_status() - except Exception as e: - raise Exception(f"Cannot conect to config server through socks5: {e}") - pass - - def __get_aws_token(self): - try: - token_url = "http://169.254.169.254/latest/api/token" - token_response = requests.put(token_url, headers={"X-aws-ec2-metadata-token-ttl-seconds": "3600"}, timeout=2) - return token_response.text - except Exception as e: - return "blank" - - def __get_current_region(self): - token = self.__get_aws_token() - metadata_url = "http://169.254.169.254/latest/dynamic/instance-identity/document" - headers = {"X-aws-ec2-metadata-token": token} - try: - response = requests.get(metadata_url, headers=headers,timeout=2) - if response.status_code == 200: - return response.json().get("region") - else: - print(f"Failed to fetch region, status code: {response.status_code}") - except Exception as e: - raise Exception(f"Region not found, are you running in EC2 environment. {e}") + subprocess.Popen(["sockd", "-d"]) def __get_secret_name_from_userdata(self): token = self.__get_aws_token() @@ -125,6 +108,23 @@ def _setup_auxilaries(self): self.__run_config_server(log_level) self.__run_socks_proxy(log_level) + + def _validate_auxilaries(self): + proxy = "socks5h://127.0.0.1:3305" + url = "http://127.0.0.1:27015/getConfig" + response = requests.get(url) + if response.status_code != 200: + raise Exception("Config server unreachable") + proxies = { + "http": proxy, + "https": proxy, + } + try: + response = requests.get(url, proxies=proxies) + response.raise_for_status() + except Exception as e: + raise Exception(f"Cannot conect to config server through socks5: {e}") + def run_compute(self): self._setup_auxilaries() self._validate_auxilaries() From 278152a42279538d61f35caab4a29e15055836c3 Mon Sep 17 00:00:00 2001 From: abuabraham-ttd Date: Mon, 18 Nov 2024 12:01:54 -0800 Subject: [PATCH 4/9] Question-1 --- scripts/aws/entrypoint.sh | 1 + 1 file changed, 1 insertion(+) diff --git a/scripts/aws/entrypoint.sh b/scripts/aws/entrypoint.sh index 37214388b..26fdcf594 100755 --- a/scripts/aws/entrypoint.sh +++ b/scripts/aws/entrypoint.sh @@ -80,6 +80,7 @@ else exit 1 fi +# DO WE NEED THIS? do we expect customers to change URL? # -- replace base URLs if both CORE_BASE_URL and OPTOUT_BASE_URL are provided # -- using hardcoded domains is fine because they should not be changed frequently if [ -n "${CORE_BASE_URL}" ] && [ "${CORE_BASE_URL}" != "null" ] && [ -n "${OPTOUT_BASE_URL}" ] && [ "${OPTOUT_BASE_URL}" != "null" ] && [ "${DEPLOYMENT_ENVIRONMENT}" != "prod" ]; then From 6929f3a3194e15e0f1ef72bd6db1f56a2e09d7c3 Mon Sep 17 00:00:00 2001 From: abuabraham-ttd Date: Mon, 18 Nov 2024 20:48:33 -0800 Subject: [PATCH 5/9] Addressing feedback --- scripts/aws/ec2.py | 227 ++++++++++++++++++-------------- scripts/aws/load_config.py | 41 ------ scripts/confidential_compute.py | 17 ++- 3 files changed, 140 insertions(+), 145 deletions(-) delete mode 100644 scripts/aws/load_config.py diff --git a/scripts/aws/ec2.py b/scripts/aws/ec2.py index bd4173953..6e34a36d1 100755 --- a/scripts/aws/ec2.py +++ b/scripts/aws/ec2.py @@ -10,170 +10,197 @@ import signal import argparse from botocore.exceptions import ClientError +from typing import Dict import sys sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -from confidential_compute import ConfidentialCompute +from confidential_compute import ConfidentialCompute, OperatorConfig + class EC2(ConfidentialCompute): def __init__(self): super().__init__() - self.config = {} + self.configs: OperatorConfig = {} - def __get_aws_token(self): + 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_response = requests.put(token_url, headers={"X-aws-ec2-metadata-token-ttl-seconds": "3600"}, timeout=2) - return token_response.text - except Exception as e: - return "blank" - - def __get_current_region(self): + response = requests.put( + 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() metadata_url = "http://169.254.169.254/latest/dynamic/instance-identity/document" headers = {"X-aws-ec2-metadata-token": token} try: - response = requests.get(metadata_url, headers=headers,timeout=2) - if response.status_code == 200: - return response.json().get("region") - else: - print(f"Failed to fetch region, status code: {response.status_code}") - except Exception as e: - raise Exception(f"Region not found, are you running in EC2 environment. {e}") + response = requests.get(metadata_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 _get_secret(self, secret_identifier): - client = boto3.client("secretsmanager", region_name=self.__get_current_region()) + def _get_secret(self, secret_identifier: str) -> Dict: + """Fetches a secret value from AWS Secrets Manager.""" + region = self.__get_current_region() + client = boto3.client("secretsmanager", region_name=region) try: secret = client.get_secret_value(SecretId=secret_identifier) return json.loads(secret["SecretString"]) except ClientError as e: - raise Exception("Unable to access secret store") - - def __add_defaults(self, configs): + raise RuntimeError(f"Unable to access Secrets Manager: {e}") + + @staticmethod + def __add_defaults(configs: Dict[str, any]) -> OperatorConfig: + """Adds default values to configuration if missing.""" configs.setdefault("enclave_memory_mb", 24576) configs.setdefault("enclave_cpu_count", 6) configs.setdefault("debug_mode", False) return configs - def __setup_vsockproxy(self, log_level): - thread_count = int((multiprocessing.cpu_count() + 1) // 2) - log_level = log_level + @staticmethod + def __error_out_on_execute(command: list, error_message: str) -> None: + """Runs a command in the background and handles exceptions.""" try: - subprocess.Popen(["/usr/bin/vsockpx", "-c", "/etc/uid2operator/proxy.yaml", "--workers", str(thread_count), "--log-level", log_level, "--daemon"]) - print("VSOCK proxy is now running in the background") - except FileNotFoundError: - print("Error: vsockpx not found. Please ensure the path is correct") + subprocess.Popen(command, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) except Exception as e: - print("Failed to start VSOCK proxy") + print(f"{error_message} \n '{' '.join(command)}': {e}") - def __run_config_server(self, log_level): + 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.__error_out_on_execute(command, "vsockpx not found. Ensure it is installed.") + + def __run_config_server(self) -> None: + """Starts the Flask configuration server.""" os.makedirs("/etc/secret/secret-value", exist_ok=True) - with open('/etc/secret/secret-value/config', 'w') as fp: - json.dump(self.configs, fp) + 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") - # TODO: Add --log-level to flask. - try: - subprocess.Popen(["./bin/flask", "run", "--host", "127.0.0.1", "--port", "27015"]) - print("Config server is now running in the background.") - except Exception as e: - print(f"Failed to start config server: {e}") + command = ["./bin/flask", "run", "--host", "127.0.0.1", "--port", "27015"] + self.__error_out_on_execute(command, "Failed to start the Flask config server.") - def __run_socks_proxy(self, log_level): - subprocess.Popen(["sockd", "-d"]) + def __run_socks_proxy(self) -> None: + """Starts the SOCKS proxy service.""" + command = ["sockd", "-d"] + self.__error_out_on_execute(command, "Failed to start socks proxy.") - def __get_secret_name_from_userdata(self): + 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_response = requests.get(user_data_url, headers={"X-aws-ec2-metadata-token": token}) - user_data = user_data_response.text - identity_scope = open("/opt/uid2operator/identity_scope.txt").read().strip() - default_name = "{}-operator-config-key".format(identity_scope.lower()) - hardcoded_value = "{}_CONFIG_SECRET_KEY".format(identity_scope.upper()) + response = requests.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_auxilaries(self): + def _setup_auxiliaries(self) -> None: + """Sets up the necessary auxiliary services and configurations.""" hostname = os.getenv("HOSTNAME", default=os.uname()[1]) - file_path = "HOSTNAME" try: - with open(file_path, "w") as file: + with open("HOSTNAME", "w") as file: file.write(hostname) - print(f"Hostname '{hostname}' written to {file_path}") + print(f"Hostname '{hostname}' written to file.") except Exception as e: - print(f"An error occurred : {e}") + """ + Ignoring error here, as we are currently not using this information anywhere. + But can be added in future for tracibility on debug + """ + print(f"Error writing hostname: {e}") + config = self._get_secret(self.__get_secret_name_from_userdata()) self.configs = self.__add_defaults(config) - log_level = 3 if self.configs['debug_mode'] else 1 + 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() - - def _validate_auxilaries(self): + def _validate_auxiliaries(self) -> None: + """Validates auxiliary services.""" proxy = "socks5h://127.0.0.1:3305" - url = "http://127.0.0.1:27015/getConfig" - response = requests.get(url) - if response.status_code != 200: - raise Exception("Config server unreachable") - proxies = { - "http": proxy, - "https": proxy, - } + config_url = "http://127.0.0.1:27015/getConfig" try: - response = requests.get(url, proxies=proxies) - response.raise_for_status() - except Exception as e: - raise Exception(f"Cannot conect to config server through socks5: {e}") + response = requests.get(config_url) + response.raise_for_status() + except requests.RequestException as e: + raise RuntimeError(f"Config server unreachable: {e}") + proxies = {"http": proxy, "https": proxy} + try: + response = requests.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}") - def run_compute(self): - self._setup_auxilaries() - self._validate_auxilaries() + def run_compute(self) -> None: + """Main execution flow for confidential compute.""" + self._setup_auxiliaries() + self._validate_auxiliaries() + self.validate_connectivity(self.configs) command = [ "nitro-cli", "run-enclave", "--eif-path", "/opt/uid2operator/uid2operator.eif", - "--memory", self.config['enclave_memory_mb'], - "--cpu-count", self.config['enclave_cpu_count'], - "--enclave-cid", 42, + "--memory", str(self.configs["enclave_memory_mb"]), + "--cpu-count", str(self.configs["enclave_cpu_count"]), + "--enclave-cid", "42", "--enclave-name", "uid2operator" ] - if self.config['debug']: - command+=["--debug-mode", "--attach-console"] + if self.configs["debug_mode"]: + command += ["--debug-mode", "--attach-console"] subprocess.run(command, check=True) - def cleanup(self): - describe_output = subprocess.check_output(["nitro-cli", "describe-enclaves"], text=True) - enclaves = json.loads(describe_output) - enclave_id = enclaves[0].get("EnclaveID") if enclaves else None - if enclave_id: - subprocess.run(["nitro-cli", "terminate-enclave", "--enclave-id", enclave_id]) - print(f"Enclave with ID {enclave_id} has been terminated.") - else: - print("No enclave found or EnclaveID is null.") - - def kill_process(self, process_name): + def cleanup(self) -> None: + """Terminates the Nitro Enclave and auxiliary processes.""" try: - result = subprocess.run( - ["pgrep", "-f", process_name], - stdout=subprocess.PIPE, - text=True, - check=False - ) + describe_output = subprocess.check_output(["nitro-cli", "describe-enclaves"], text=True) + enclaves = json.loads(describe_output) + enclave_id = enclaves[0].get("EnclaveID") if enclaves else None + if enclave_id: + subprocess.run(["nitro-cli", "terminate-enclave", "--enclave-id", enclave_id]) + print(f"Terminated enclave with ID: {enclave_id}") + else: + print("No active enclaves found.") + except subprocess.SubprocessError as e: + raise (f"Error during cleanup: {e}") + + def kill_process(self, process_name: str) -> None: + """Kills a process by its name.""" + try: + 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"{process_name} exited") + print(f"Killed process '{process_name}'.") else: - print(f"Process {process_name} not found") + print(f"No process named '{process_name}' found.") except Exception as e: - print(f"Failed to shut down {process_name}: {e}") + print(f"Error killing process '{process_name}': {e}") + if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("-o", "--operation", required=False) + 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 and args.operation == "stop": + if args.operation == "stop": ec2.cleanup() - [ec2.kill_process(process) for process in ["vsockpx", "sockd", "vsock-proxy", "nohup"]] + for process in ["vsockpx", "sockd", "vsock-proxy"]: + ec2.kill_process(process) else: ec2.run_compute() + \ No newline at end of file diff --git a/scripts/aws/load_config.py b/scripts/aws/load_config.py deleted file mode 100644 index 9f0446a49..000000000 --- a/scripts/aws/load_config.py +++ /dev/null @@ -1,41 +0,0 @@ -import os -import boto3 -import base64 -import json -from botocore.exceptions import ClientError - -secret_name = os.environ['UID2_CONFIG_SECRET_KEY'] -region_name = os.environ['AWS_REGION_NAME'] -aws_access_key_id = os.environ['AWS_ACCESS_KEY_ID'] -secret_key = os.environ['AWS_SECRET_KEY'] -session_token = os.environ['AWS_SESSION_TOKEN'] - -def get_secret(): - session = boto3.session.Session() - client = session.client( - service_name='secretsmanager', - region_name=region_name, - aws_access_key_id = aws_access_key_id, - aws_secret_access_key = secret_key, - aws_session_token = session_token - ) - try: - get_secret_value_response = client.get_secret_value( - SecretId=secret_name - ) - except ClientError as e: - raise e - else: - if 'SecretString' in get_secret_value_response: - secret = get_secret_value_response['SecretString'] - else: - decoded_binary_secret = base64.b64decode(get_secret_value_response['SecretBinary']) - - return secret - -def get_config(): - result = get_secret() - conf = json.loads(result) - print(result) - -get_config() diff --git a/scripts/confidential_compute.py b/scripts/confidential_compute.py index a60220bcd..8469ea63d 100644 --- a/scripts/confidential_compute.py +++ b/scripts/confidential_compute.py @@ -3,6 +3,15 @@ import socket from urllib.parse import urlparse from abc import ABC, abstractmethod +from typing import TypedDict, Dict + +class OperatorConfig(TypedDict): + enclave_memory_mb: int + enclave_cpu_count: int + debug_mode: bool + api_token: str + core_base_url: str + optout_base_url: str class ConfidentialCompute(ABC): @@ -16,7 +25,7 @@ def _get_secret(self, secret_identifier): """ pass - def validate_operator_key(self, secrets): + def validate_operator_key(self, secrets: OperatorConfig): """ Validates operator key if following new pattern. Ignores otherwise """ @@ -31,7 +40,7 @@ def validate_operator_key(self, secrets): raise Exception("Operator key does not match the environment") return True - def validate_connectivity(self, config): + def validate_connectivity(self, config: OperatorConfig): """ Validates core/optout is accessible. """ @@ -52,14 +61,14 @@ def validate_connectivity(self, config): return @abstractmethod - def _setup_auxilaries(self, secrets): + def _setup_auxiliaries(self): """ Sets up auxilary processes required to confidential compute """ pass @abstractmethod - def _validate_auxilaries(self, secrets): + def _validate_auxiliaries(self): """ Validates auxilary services are running """ From 1daba547ade63dbfb291004cb03cb5b2935f82bb Mon Sep 17 00:00:00 2001 From: abuabraham-ttd Date: Mon, 18 Nov 2024 20:59:36 -0800 Subject: [PATCH 6/9] Addressing feedback-2 --- scripts/confidential_compute.py | 96 ++++++++++++++++----------------- 1 file changed, 48 insertions(+), 48 deletions(-) diff --git a/scripts/confidential_compute.py b/scripts/confidential_compute.py index 8469ea63d..32e509ad3 100644 --- a/scripts/confidential_compute.py +++ b/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 \ No newline at end of file From bdfe34bfe8348ff5801fc4a5634c0e7cc4796186 Mon Sep 17 00:00:00 2001 From: abuabraham-ttd Date: Mon, 18 Nov 2024 21:40:47 -0800 Subject: [PATCH 7/9] Add default core, optout base url --- scripts/aws/ec2.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/scripts/aws/ec2.py b/scripts/aws/ec2.py index 6e34a36d1..e60544dbf 100755 --- a/scripts/aws/ec2.py +++ b/scripts/aws/ec2.py @@ -61,6 +61,8 @@ def __add_defaults(configs: Dict[str, any]) -> OperatorConfig: configs.setdefault("enclave_memory_mb", 24576) configs.setdefault("enclave_cpu_count", 6) configs.setdefault("debug_mode", False) + configs.setdefault("core_base_url", "https://core.uidapi.com" if configs["environment"] == "prod" else "https://core-integ.uidapi.com") + configs.setdefault("optout_base_url", "https://optout.uidapi.com" if configs["environment"] == "prod" else "https://optout-integ.uidapi.com") return configs @staticmethod From 01c38503c36b20c6912b9d1646137578a493c66d Mon Sep 17 00:00:00 2001 From: abuabraham-ttd Date: Tue, 19 Nov 2024 19:11:28 -0800 Subject: [PATCH 8/9] Add sample Azure implementation --- scripts/aws/ec2.py | 46 ++++++----- scripts/azure-cc/Dockerfile | 17 +++- scripts/azure-cc/conf/azure.py | 136 ++++++++++++++++++++++++++++++++ scripts/confidential_compute.py | 1 + 4 files changed, 177 insertions(+), 23 deletions(-) create mode 100644 scripts/azure-cc/conf/azure.py diff --git a/scripts/aws/ec2.py b/scripts/aws/ec2.py index e60544dbf..11578dc9a 100755 --- a/scripts/aws/ec2.py +++ b/scripts/aws/ec2.py @@ -51,6 +51,7 @@ def _get_secret(self, secret_identifier: str) -> Dict: client = boto3.client("secretsmanager", region_name=region) try: secret = client.get_secret_value(SecretId=secret_identifier) + #TODO: validate secret string has operator key, environment and other required values return json.loads(secret["SecretString"]) except ClientError as e: raise RuntimeError(f"Unable to access Secrets Manager: {e}") @@ -65,37 +66,42 @@ def __add_defaults(configs: Dict[str, any]) -> OperatorConfig: configs.setdefault("optout_base_url", "https://optout.uidapi.com" if configs["environment"] == "prod" else "https://optout-integ.uidapi.com") return configs - @staticmethod - def __error_out_on_execute(command: list, error_message: str) -> None: - """Runs a command in the background and handles exceptions.""" - try: - subprocess.Popen(command, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) - except Exception as e: - print(f"{error_message} \n '{' '.join(command)}': {e}") - def __setup_vsockproxy(self, log_level: int) -> None: - """Sets up the vsock proxy service.""" + """ + Sets up the vsock proxy service. + TODO: Evaluate adding vsock logging based on log_level here + """ 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.__error_out_on_execute(command, "vsockpx not found. Ensure it is installed.") + subprocess.run(command) - def __run_config_server(self) -> None: - """Starts the Flask configuration server.""" + def __run_config_server(self,log_level) -> None: + """ + Starts the Flask configuration server. + TODO: Based on log level add logging to flask + """ 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", "127.0.0.1", "--port", "27015"] - self.__error_out_on_execute(command, "Failed to start the Flask config server.") - - def __run_socks_proxy(self) -> None: - """Starts the SOCKS proxy service.""" + try: + subprocess.Popen(command, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + except Exception as e: + print(f"Failed to start the Flask config server.\n '{' '.join(command)}': {e}") + raise RuntimeError ("Failed to start required flask server") + + def __run_socks_proxy(self, log_level) -> None: + """ + Starts the SOCKS proxy service. + TODO: Based on log level add logging to sockd + """ command = ["sockd", "-d"] - self.__error_out_on_execute(command, "Failed to start socks proxy.") + subprocess.run(command) def __get_secret_name_from_userdata(self) -> str: """Extracts the secret name from EC2 user data.""" @@ -125,13 +131,13 @@ def _setup_auxiliaries(self) -> None: But can be added in future for tracibility on debug """ print(f"Error writing hostname: {e}") - + config = self._get_secret(self.__get_secret_name_from_userdata()) self.configs = self.__add_defaults(config) log_level = 3 if self.configs["debug_mode"] else 1 self.__setup_vsockproxy(log_level) - self.__run_config_server() - self.__run_socks_proxy() + self.__run_config_server(log_level) + self.__run_socks_proxy(log_level) def _validate_auxiliaries(self) -> None: """Validates auxiliary services.""" diff --git a/scripts/azure-cc/Dockerfile b/scripts/azure-cc/Dockerfile index bb0c96b70..5704f15fc 100644 --- a/scripts/azure-cc/Dockerfile +++ b/scripts/azure-cc/Dockerfile @@ -17,6 +17,17 @@ ENV IMAGE_VERSION=${IMAGE_VERSION} ENV REGION=default ENV LOKI_HOSTNAME=loki +RUN apt-get update && apt-get install -y --no-install-recommends \ + python3 \ + python3-pip \ + && apt-get clean && rm -rf /var/lib/apt/lists/* + +RUN python3 -m pip install --upgrade pip + +RUN pip install --no-cache-dir \ + azure-identity \ + azure-keyvault-secrets + COPY ./target/${JAR_NAME}-${JAR_VERSION}-jar-with-dependencies.jar /app/${JAR_NAME}-${JAR_VERSION}.jar COPY ./target/${JAR_NAME}-${JAR_VERSION}-sources.jar /app COPY ./target/${JAR_NAME}-${JAR_VERSION}-static.tar.gz /app/static.tar.gz @@ -25,10 +36,10 @@ COPY ./conf/*.xml /app/conf/ RUN tar xzvf /app/static.tar.gz --no-same-owner --no-same-permissions && rm -f /app/static.tar.gz -COPY ./entrypoint.sh /app/ -RUN chmod a+x /app/entrypoint.sh +COPY ./azure.py /app/ +RUN chmod a+x /app/azure.py RUN adduser -D uid2-operator && mkdir -p /opt/uid2 && chmod 777 -R /opt/uid2 && mkdir -p /app && chmod 705 -R /app && mkdir -p /app/file-uploads && chmod 777 -R /app/file-uploads USER uid2-operator -CMD ["/app/entrypoint.sh"] +CMD ["/app/azure.py"] diff --git a/scripts/azure-cc/conf/azure.py b/scripts/azure-cc/conf/azure.py new file mode 100644 index 000000000..44d5e80be --- /dev/null +++ b/scripts/azure-cc/conf/azure.py @@ -0,0 +1,136 @@ +import os +import subprocess +import time +import json +import sys +import requests +import re +from typing import Dict +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from confidential_compute import ConfidentialCompute, OperatorConfig +from azure.identity import DefaultAzureCredential +from azure.keyvault.secrets import SecretClient + +class AzureCC(ConfidentialCompute): + + def __init__(self): + super().__init__() + self.configs: OperatorConfig = {} + + + def _get_secret(self, secret_identifier): + """Fetches a secret value from Azure Key Value, reads environment variables and returns config""" + key_vault_url = "https://{}.vault.azure.net/".format(secret_identifier["key_vault"]) + credential = DefaultAzureCredential() + secret_client = SecretClient(vault_url=key_vault_url, credential=credential) + try: + config = { + "api_key" : secret_client.get_secret(secret_identifier["secret_name"]), + "environment": os.getenv("DEPLOYMENT_ENVIRONMENT"), + "core_base_url": os.getenv("CORE_BASE_URL"), + "optout_base_url": os.getenv("OPTOUT_BASE_URL") + } + return {key: value for key, value in config.items() if value is not None} + except Exception as e: + raise RuntimeError(f"Unable to access Secrets Manager: {e}") + + def _setup_auxiliaries(self, secrets): + """Sets up auxiliary configurations (placeholder for extension).""" + pass + + def __validate_sidecar(self): + """Validates the required sidecar is running""" + url = "http://169.254.169.254/ping" + delay = 1 + max_retries = 15 + while True: + try: + response = requests.get(url, timeout=5) + if response.status_code == 200: + print("Sidecar started") + break + except requests.RequestException: + print(f"Sidecar not started. Retrying in {delay} seconds...") + time.sleep(delay) + if delay > max_retries: + raise RuntimeError("Unable to start operator as sidecar failed to start") + delay += 1 + + + def _validate_auxiliaries(self, secrets): + """Validates the presence of required environment variables, and sidecar is up""" + self.__validate_sidecar() + config_env_vars = [ + "VAULT_NAME", + "OPERATOR_KEY_SECRET_NAME", + "DEPLOYMENT_ENVIRONMENT" + ] + pre_set_env_vars = [ + "JAR_NAME", + "JAR_VERSION" + ] + for variable in (config_env_vars + pre_set_env_vars): + value = os.getenv(variable) + if not value: + raise ValueError("{} is not set. Please update it".format(variable)) + if os.getenv("DEPLOYMENT_ENVIRONMENT") not in ["prod","integ"]: + raise ValueError("DEPLOYMENT_ENVIRONMENT should be prod/integ. It is currently set to {}".format(os.getenv("DEPLOYMENT_ENVIRONMENT"))) + + @staticmethod + def __add_defaults(configs: Dict[str, any]) -> OperatorConfig: + """Adds default values to configuration if missing.""" + configs.setdefault("enclave_memory_mb", -1) + configs.setdefault("enclave_cpu_count", -1) + configs.setdefault("debug_mode", False) + configs.setdefault("core_base_url", "https://core.uidapi.com" if configs["environment"] == "prod" else "https://core-integ.uidapi.com") + configs.setdefault("optout_base_url", "https://optout.uidapi.com" if configs["environment"] == "prod" else "https://optout-integ.uidapi.com") + return configs + + def __update_config_file(self, config_path): + """Updates configuration file with base URLs if in a non-production environment.""" + if not os.path.exists(config_path): + raise FileNotFoundError(f"Configuration file not found: {config_path}") + with open(config_path) as f: + config_data = json.load(f) + if all([os.getenv("CORE_BASE_URL"), os.getenv("OPTOUT_BASE_URL")]) and self.configs["environment"] != "prod": + config_data = re.sub(r"https://core-integ\.uidapi\.com", os.getenv("CORE_BASE_URL"), config_data) + config_data = re.sub(r"https://optout-integ\.uidapi\.com", os.getenv("OPTOUT_BASE_URL"), config_data) + with open(config_path, "w") as file: + file.write(config_data) + + + def run_compute(self): + """Main execution flow for confidential compute.""" + self._setup_auxiliaries(None) + self._validate_auxiliaries(None) + secret_identifier = { + "key_vault": os.getenv("OPERATOR_KEY_SECRET_NAME"), + "secret_name": os.getenv("VAULT_NAME") + } + self.configs = self.__add_defaults(self._get_secret(secret_identifier)) + self.validate_operator_key(self.configs) + self.validate_connectivity(self.configs) + os.environ["azure_vault_name"] = os.getenv("VAULT_NAME") + os.environ["azure_secret_name"] = os.getenv("OPERATOR_KEY_SECRET_NAME") + config_path="/app/conf/${}-uid2-config.json".format(os.getenv("DEPLOYMENT_ENVIRONMENT")) + self.__update_config_file(config_path=config_path) + java_command = [ + "java", + "-XX:MaxRAMPercentage=95", + "-XX:-UseCompressedOops", + "-XX:+PrintFlagsFinal", + "-Djava.security.egd=file:/dev/./urandom", + "-Dvertx.logger-delegate-factory-class-name=io.vertx.core.logging.SLF4JLogDelegateFactory", + "-Dlogback.configurationFile=/app/conf/logback.xml", + "-Dvertx-config-path={}".format(config_path), + "-jar", + "{}-{}.jar".format(os.getenv("JAR_NAME"), os.getenv("JAR_VERSION")) + ] + try: + subprocess.run(java_command, check=True) + except subprocess.CalledProcessError as e: + print(f"Error starting the Java application: {e}") + + +if __name__ == "__main__": + AzureCC().run_compute() \ No newline at end of file diff --git a/scripts/confidential_compute.py b/scripts/confidential_compute.py index 32e509ad3..124f2b7fc 100644 --- a/scripts/confidential_compute.py +++ b/scripts/confidential_compute.py @@ -13,6 +13,7 @@ class OperatorConfig(TypedDict): api_token: str core_base_url: str optout_base_url: str + environment: str class ConfidentialCompute(ABC): From d8e2663620ec9daf533c12ad5a27aec173051a48 Mon Sep 17 00:00:00 2001 From: abuabraham-ttd Date: Sat, 23 Nov 2024 16:54:37 -0800 Subject: [PATCH 9/9] commit local --- .github/actions/build_aws_eif/action.yaml | 4 +- scripts/aws/EUID_CloudFormation.template.yml | 2 +- scripts/aws/UID_CloudFormation.template.yml | 2 +- scripts/aws/ec2.py | 44 +++++++++---------- .../uid2-operator-ami/ansible/playbook.yml | 6 +++ scripts/azure-cc/conf/azure.py | 27 ++++++------ scripts/confidential_compute.py | 40 +++++++++++------ 7 files changed, 70 insertions(+), 55 deletions(-) diff --git a/.github/actions/build_aws_eif/action.yaml b/.github/actions/build_aws_eif/action.yaml index f17523a44..3d08f3942 100644 --- a/.github/actions/build_aws_eif/action.yaml +++ b/.github/actions/build_aws_eif/action.yaml @@ -96,8 +96,8 @@ 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/confidential_compute.py ${ARTIFACTS_OUTPUT_DIR}/ + cp ./scripts/aws/ec2.py ${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/scripts/aws/EUID_CloudFormation.template.yml b/scripts/aws/EUID_CloudFormation.template.yml index 9c5982488..519665e26 100644 --- a/scripts/aws/EUID_CloudFormation.template.yml +++ b/scripts/aws/EUID_CloudFormation.template.yml @@ -155,7 +155,7 @@ Resources: KmsKeyId: !GetAtt KMSKey.Arn Name: !Sub 'euid-config-stack-${AWS::StackName}' SecretString: !Sub '{ - "api_token":"${APIToken}", + "operator_key":"${APIToken}", "service_instances":6, "enclave_cpu_count":6, "enclave_memory_mb":24576, diff --git a/scripts/aws/UID_CloudFormation.template.yml b/scripts/aws/UID_CloudFormation.template.yml index 711d1ab0e..468afa7a7 100644 --- a/scripts/aws/UID_CloudFormation.template.yml +++ b/scripts/aws/UID_CloudFormation.template.yml @@ -183,7 +183,7 @@ Resources: KmsKeyId: !GetAtt KMSKey.Arn Name: !Sub 'uid2-config-stack-${AWS::StackName}' SecretString: !Sub '{ - "api_token":"${APIToken}", + "operator_key":"${APIToken}", "service_instances":6, "enclave_cpu_count":6, "enclave_memory_mb":24576, diff --git a/scripts/aws/ec2.py b/scripts/aws/ec2.py index 11578dc9a..e6b29cb62 100755 --- a/scripts/aws/ec2.py +++ b/scripts/aws/ec2.py @@ -6,21 +6,21 @@ import subprocess import re import multiprocessing -import requests +import requests #need requests[socks] import signal import argparse from botocore.exceptions import ClientError from typing import Dict import sys sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -from confidential_compute import ConfidentialCompute, OperatorConfig +from confidential_compute import ConfidentialCompute, ConfidentialComputeConfig class EC2(ConfidentialCompute): def __init__(self): super().__init__() - self.configs: OperatorConfig = {} + self.configs: ConfidentialComputeConfig = {} def __get_aws_token(self) -> str: """Fetches a temporary AWS EC2 metadata token.""" @@ -45,19 +45,19 @@ def __get_current_region(self) -> str: except requests.RequestException as e: raise RuntimeError(f"Failed to fetch region: {e}") - def _get_secret(self, secret_identifier: str) -> Dict: + def _get_secret(self, secret_identifier: str) -> ConfidentialComputeConfig: + secret_identifier = "uid2-config-stack-tjm-unvalidate-eif-test1" """Fetches a secret value from AWS Secrets Manager.""" region = self.__get_current_region() client = boto3.client("secretsmanager", region_name=region) try: secret = client.get_secret_value(SecretId=secret_identifier) - #TODO: validate secret string has operator key, environment and other required values - return json.loads(secret["SecretString"]) + return self.__add_defaults(json.loads(secret["SecretString"])) except ClientError as e: - raise RuntimeError(f"Unable to access Secrets Manager: {e}") + raise RuntimeError(f"Unable to access Secrets Manager {secret_identifier}: {e}") @staticmethod - def __add_defaults(configs: Dict[str, any]) -> OperatorConfig: + def __add_defaults(configs: Dict[str, any]) -> ConfidentialComputeConfig: """Adds default values to configuration if missing.""" configs.setdefault("enclave_memory_mb", 24576) configs.setdefault("enclave_cpu_count", 6) @@ -100,7 +100,7 @@ def __run_socks_proxy(self, log_level) -> None: Starts the SOCKS proxy service. TODO: Based on log level add logging to sockd """ - command = ["sockd", "-d"] + command = ["sockd", "-D"] subprocess.run(command) def __get_secret_name_from_userdata(self) -> str: @@ -131,9 +131,7 @@ def _setup_auxiliaries(self) -> None: But can be added in future for tracibility on debug """ print(f"Error writing hostname: {e}") - - config = self._get_secret(self.__get_secret_name_from_userdata()) - self.configs = self.__add_defaults(config) + 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) @@ -141,7 +139,7 @@ def _setup_auxiliaries(self) -> None: def _validate_auxiliaries(self) -> None: """Validates auxiliary services.""" - proxy = "socks5h://127.0.0.1:3305" + proxy = "socks5://127.0.0.1:3306" config_url = "http://127.0.0.1:27015/getConfig" try: response = requests.get(config_url) @@ -183,19 +181,21 @@ def cleanup(self) -> None: print(f"Terminated enclave with ID: {enclave_id}") else: print("No active enclaves found.") + self.__kill_auxiliaries() except subprocess.SubprocessError as e: raise (f"Error during cleanup: {e}") - def kill_process(self, process_name: str) -> None: + def __kill_auxiliaries(self) -> None: """Kills a process by its name.""" try: - 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.") + for process_name in ["vsockpx", "sockd"]: + 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}") @@ -207,8 +207,6 @@ def kill_process(self, process_name: str) -> None: ec2 = EC2() if args.operation == "stop": ec2.cleanup() - for process in ["vsockpx", "sockd", "vsock-proxy"]: - ec2.kill_process(process) else: ec2.run_compute() \ 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 3629b34e3..4a8d302d3 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 confidential_compute script + ansible.builtin.copy: + src: /tmp/artifacts/confidential_compute.py + dest: /opt/uid2operator/confidential_compute.py + remote_src: yes + - name: Install starter script ansible.builtin.copy: src: /tmp/artifacts/ec2.py diff --git a/scripts/azure-cc/conf/azure.py b/scripts/azure-cc/conf/azure.py index 44d5e80be..57b1cabed 100644 --- a/scripts/azure-cc/conf/azure.py +++ b/scripts/azure-cc/conf/azure.py @@ -7,7 +7,7 @@ import re from typing import Dict sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -from confidential_compute import ConfidentialCompute, OperatorConfig +from confidential_compute import ConfidentialCompute, ConfidentialComputeConfig, OperatorConfig from azure.identity import DefaultAzureCredential from azure.keyvault.secrets import SecretClient @@ -15,10 +15,9 @@ class AzureCC(ConfidentialCompute): def __init__(self): super().__init__() - self.configs: OperatorConfig = {} + self.configs: ConfidentialComputeConfig = {} - - def _get_secret(self, secret_identifier): + def _get_secret(self, secret_identifier) -> ConfidentialComputeConfig: """Fetches a secret value from Azure Key Value, reads environment variables and returns config""" key_vault_url = "https://{}.vault.azure.net/".format(secret_identifier["key_vault"]) credential = DefaultAzureCredential() @@ -30,7 +29,7 @@ def _get_secret(self, secret_identifier): "core_base_url": os.getenv("CORE_BASE_URL"), "optout_base_url": os.getenv("OPTOUT_BASE_URL") } - return {key: value for key, value in config.items() if value is not None} + return self.__add_defaults({key: value for key, value in config.items() if value is not None}) except Exception as e: raise RuntimeError(f"Unable to access Secrets Manager: {e}") @@ -77,7 +76,7 @@ def _validate_auxiliaries(self, secrets): raise ValueError("DEPLOYMENT_ENVIRONMENT should be prod/integ. It is currently set to {}".format(os.getenv("DEPLOYMENT_ENVIRONMENT"))) @staticmethod - def __add_defaults(configs: Dict[str, any]) -> OperatorConfig: + def __add_defaults(configs: Dict[str, any]) -> ConfidentialComputeConfig: """Adds default values to configuration if missing.""" configs.setdefault("enclave_memory_mb", -1) configs.setdefault("enclave_cpu_count", -1) @@ -85,9 +84,10 @@ def __add_defaults(configs: Dict[str, any]) -> OperatorConfig: configs.setdefault("core_base_url", "https://core.uidapi.com" if configs["environment"] == "prod" else "https://core-integ.uidapi.com") configs.setdefault("optout_base_url", "https://optout.uidapi.com" if configs["environment"] == "prod" else "https://optout-integ.uidapi.com") return configs - - def __update_config_file(self, config_path): - """Updates configuration file with base URLs if in a non-production environment.""" + + #TODO: This is repeated in GCP, EC2 + def __get_overriden_configs(self, config_path) -> OperatorConfig: + """Returns the required configurations for operator. Only overrides if environment is integ""" if not os.path.exists(config_path): raise FileNotFoundError(f"Configuration file not found: {config_path}") with open(config_path) as f: @@ -95,9 +95,7 @@ def __update_config_file(self, config_path): if all([os.getenv("CORE_BASE_URL"), os.getenv("OPTOUT_BASE_URL")]) and self.configs["environment"] != "prod": config_data = re.sub(r"https://core-integ\.uidapi\.com", os.getenv("CORE_BASE_URL"), config_data) config_data = re.sub(r"https://optout-integ\.uidapi\.com", os.getenv("OPTOUT_BASE_URL"), config_data) - with open(config_path, "w") as file: - file.write(config_data) - + return config_data def run_compute(self): """Main execution flow for confidential compute.""" @@ -107,13 +105,14 @@ def run_compute(self): "key_vault": os.getenv("OPERATOR_KEY_SECRET_NAME"), "secret_name": os.getenv("VAULT_NAME") } - self.configs = self.__add_defaults(self._get_secret(secret_identifier)) + self.configs = self._get_secret(secret_identifier) self.validate_operator_key(self.configs) self.validate_connectivity(self.configs) os.environ["azure_vault_name"] = os.getenv("VAULT_NAME") os.environ["azure_secret_name"] = os.getenv("OPERATOR_KEY_SECRET_NAME") config_path="/app/conf/${}-uid2-config.json".format(os.getenv("DEPLOYMENT_ENVIRONMENT")) - self.__update_config_file(config_path=config_path) + with open(config_path, "w") as file: + file.write(self.__get_overriden_configs(config_path=config_path)) java_command = [ "java", "-XX:MaxRAMPercentage=95", diff --git a/scripts/confidential_compute.py b/scripts/confidential_compute.py index 124f2b7fc..c2fa44055 100644 --- a/scripts/confidential_compute.py +++ b/scripts/confidential_compute.py @@ -3,22 +3,36 @@ import socket from urllib.parse import urlparse from abc import ABC, abstractmethod -from typing import TypedDict +from typing import TypedDict, NotRequired -class OperatorConfig(TypedDict): +class ConfidentialComputeConfig(TypedDict): enclave_memory_mb: int enclave_cpu_count: int debug_mode: bool - api_token: str + operator_key: str core_base_url: str optout_base_url: str environment: str - +class OperatorConfig(TypedDict): + sites_metadata_path: str + clients_metadata_path: str + keysets_metadata_path: str + keyset_keys_metadata_path: str + salts_metadata_path: str + services_metadata_path: str + service_links_metadata_path: str + optout_metadata_path: str + core_attest_url: str + optout_api_uri: str + optout_s3_folder: str + identity_token_expires_after_seconds: str + client_side_keypairs_metadata_path: NotRequired[str] + class ConfidentialCompute(ABC): @abstractmethod - def _get_secret(self, secret_identifier: str) -> OperatorConfig: + def _get_secret(self, secret_identifier: str) -> ConfidentialComputeConfig: """ Fetches the secret from a secret store. @@ -27,18 +41,17 @@ def _get_secret(self, secret_identifier: str) -> OperatorConfig: """ pass - def validate_operator_key(self, secrets: OperatorConfig) -> bool: + def validate_operator_key(self, secrets: ConfidentialComputeConfig) -> bool: """ Validates the operator key format and its environment alignment.""" - api_token = secrets.get("api_token") - if not api_token: + operator_key = secrets.get("operator_key") + if not operator_key: raise ValueError("API token is missing from the configuration.") - pattern = r"^(UID2|EUID)-.\-(I|P)-\d+-\*$" - if re.match(pattern, api_token): + if re.match(pattern, operator_key): 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: + if operator_key.split("-")[2] != expected_env: raise ValueError( f"Operator key does not match the expected environment ({expected_env})." ) @@ -50,7 +63,7 @@ def __resolve_hostname(url: str) -> str: hostname = urlparse(url).netloc return socket.gethostbyname(hostname) - def validate_connectivity(self, config: OperatorConfig) -> None: + def validate_connectivity(self, config: ConfidentialComputeConfig) -> None: """ Validates that the core and opt-out URLs are accessible.""" try: core_url = config["core_base_url"] @@ -59,7 +72,6 @@ def validate_connectivity(self, config: OperatorConfig) -> None: 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." @@ -80,4 +92,4 @@ def _validate_auxiliaries(self) -> None: @abstractmethod def run_compute(self) -> None: """ Runs confidential computing.""" - pass \ No newline at end of file + pass