Skip to content

Commit

Permalink
commit local
Browse files Browse the repository at this point in the history
  • Loading branch information
abuabraham-ttd committed Nov 24, 2024
1 parent 01c3850 commit d8e2663
Show file tree
Hide file tree
Showing 7 changed files with 70 additions and 55 deletions.
4 changes: 2 additions & 2 deletions .github/actions/build_aws_eif/action.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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}/
Expand Down
2 changes: 1 addition & 1 deletion scripts/aws/EUID_CloudFormation.template.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
2 changes: 1 addition & 1 deletion scripts/aws/UID_CloudFormation.template.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
44 changes: 21 additions & 23 deletions scripts/aws/ec2.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand All @@ -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)
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -131,17 +131,15 @@ 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)
self.__run_socks_proxy(log_level)

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)
Expand Down Expand Up @@ -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}")

Expand All @@ -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()

6 changes: 6 additions & 0 deletions scripts/aws/uid2-operator-ami/ansible/playbook.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
27 changes: 13 additions & 14 deletions scripts/azure-cc/conf/azure.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,18 +7,17 @@
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

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()
Expand All @@ -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}")

Expand Down Expand Up @@ -77,27 +76,26 @@ 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)
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."""

#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:
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)

return config_data

def run_compute(self):
"""Main execution flow for confidential compute."""
Expand All @@ -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",
Expand Down
40 changes: 26 additions & 14 deletions scripts/confidential_compute.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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})."
)
Expand All @@ -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"]
Expand All @@ -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."
Expand All @@ -80,4 +92,4 @@ def _validate_auxiliaries(self) -> None:
@abstractmethod
def run_compute(self) -> None:
""" Runs confidential computing."""
pass
pass

0 comments on commit d8e2663

Please sign in to comment.