|
| 1 | +#!/usr/bin/env python3 |
| 2 | + |
| 3 | +import boto3 |
| 4 | +import json |
| 5 | +import os |
| 6 | +import subprocess |
| 7 | +import re |
| 8 | +import multiprocessing |
| 9 | +import requests |
| 10 | +import signal |
| 11 | +import argparse |
| 12 | +from botocore.exceptions import ClientError |
| 13 | +import sys |
| 14 | +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) |
| 15 | +from confidential_compute import ConfidentialCompute |
| 16 | + |
| 17 | +class EC2(ConfidentialCompute): |
| 18 | + |
| 19 | + def __init__(self): |
| 20 | + super().__init__() |
| 21 | + self.config = {} |
| 22 | + |
| 23 | + def _get_secret(self, secret_identifier): |
| 24 | + client = boto3.client("secretsmanager", region_name=self.__get_current_region()) |
| 25 | + try: |
| 26 | + secret = client.get_secret_value(SecretId=secret_identifier) |
| 27 | + return json.loads(secret["SecretString"]) |
| 28 | + except ClientError as e: |
| 29 | + raise Exception("Unable to access secret store") |
| 30 | + |
| 31 | + def __add_defaults(self, configs): |
| 32 | + configs.setdefault("enclave_memory_mb", 24576) |
| 33 | + configs.setdefault("enclave_cpu_count", 6) |
| 34 | + configs.setdefault("debug_mode", False) |
| 35 | + return configs |
| 36 | + |
| 37 | + def __setup_vsockproxy(self, log_level): |
| 38 | + thread_count = int((multiprocessing.cpu_count() + 1) // 2) |
| 39 | + log_level = log_level |
| 40 | + try: |
| 41 | + subprocess.Popen(["/usr/bin/vsockpx", "-c", "/etc/uid2operator/proxy.yaml", "--workers", str(thread_count), "--log-level", log_level, "--daemon"]) |
| 42 | + print("VSOCK proxy is now running in the background") |
| 43 | + except FileNotFoundError: |
| 44 | + print("Error: vsockpx not found. Please ensure the path is correct") |
| 45 | + except Exception as e: |
| 46 | + print("Failed to start VSOCK proxy") |
| 47 | + |
| 48 | + def __run_config_server(self, log_level): |
| 49 | + os.makedirs("/etc/secret/secret-value", exist_ok=True) |
| 50 | + with open('/etc/secret/secret-value/config', 'w') as fp: |
| 51 | + json.dump(self.configs, fp) |
| 52 | + os.chdir("/opt/uid2operator/config-server") |
| 53 | + # TODO: Add --log-level to flask. |
| 54 | + try: |
| 55 | + subprocess.Popen(["./bin/flask", "run", "--host", "127.0.0.1", "--port", "27015"]) |
| 56 | + print("Config server is now running in the background.") |
| 57 | + except Exception as e: |
| 58 | + print(f"Failed to start config server: {e}") |
| 59 | + |
| 60 | + def __run_socks_proxy(self, log_level): |
| 61 | + subprocess.Popen(["sockd", "-d"]) |
| 62 | + |
| 63 | + def _validate_auxilaries(self): |
| 64 | + proxy = "socks5h://127.0.0.1:3305" |
| 65 | + url = "http://127.0.0.1:27015/getConfig" |
| 66 | + response = requests.get(url) |
| 67 | + if response.status_code != 200: |
| 68 | + raise Exception("Config server unreachable") |
| 69 | + proxies = { |
| 70 | + "http": proxy, |
| 71 | + "https": proxy, |
| 72 | + } |
| 73 | + try: |
| 74 | + response = requests.get(url, proxies=proxies) |
| 75 | + response.raise_for_status() |
| 76 | + except Exception as e: |
| 77 | + raise Exception(f"Cannot conect to config server through socks5: {e}") |
| 78 | + pass |
| 79 | + |
| 80 | + def __get_aws_token(self): |
| 81 | + try: |
| 82 | + token_url = "http://169.254.169.254/latest/api/token" |
| 83 | + token_response = requests.put(token_url, headers={"X-aws-ec2-metadata-token-ttl-seconds": "3600"}, timeout=2) |
| 84 | + return token_response.text |
| 85 | + except Exception as e: |
| 86 | + return "blank" |
| 87 | + |
| 88 | + def __get_current_region(self): |
| 89 | + token = self.__get_aws_token() |
| 90 | + metadata_url = "http://169.254.169.254/latest/dynamic/instance-identity/document" |
| 91 | + headers = {"X-aws-ec2-metadata-token": token} |
| 92 | + try: |
| 93 | + response = requests.get(metadata_url, headers=headers,timeout=2) |
| 94 | + if response.status_code == 200: |
| 95 | + return response.json().get("region") |
| 96 | + else: |
| 97 | + print(f"Failed to fetch region, status code: {response.status_code}") |
| 98 | + except Exception as e: |
| 99 | + raise Exception(f"Region not found, are you running in EC2 environment. {e}") |
| 100 | + |
| 101 | + def __get_secret_name_from_userdata(self): |
| 102 | + token = self.__get_aws_token() |
| 103 | + user_data_url = "http://169.254.169.254/latest/user-data" |
| 104 | + user_data_response = requests.get(user_data_url, headers={"X-aws-ec2-metadata-token": token}) |
| 105 | + user_data = user_data_response.text |
| 106 | + identity_scope = open("/opt/uid2operator/identity_scope.txt").read().strip() |
| 107 | + default_name = "{}-operator-config-key".format(identity_scope.lower()) |
| 108 | + hardcoded_value = "{}_CONFIG_SECRET_KEY".format(identity_scope.upper()) |
| 109 | + match = re.search(rf'^export {hardcoded_value}="(.+?)"$', user_data, re.MULTILINE) |
| 110 | + return match.group(1) if match else default_name |
| 111 | + |
| 112 | + def _setup_auxilaries(self): |
| 113 | + hostname = os.getenv("HOSTNAME", default=os.uname()[1]) |
| 114 | + file_path = "HOSTNAME" |
| 115 | + try: |
| 116 | + with open(file_path, "w") as file: |
| 117 | + file.write(hostname) |
| 118 | + print(f"Hostname '{hostname}' written to {file_path}") |
| 119 | + except Exception as e: |
| 120 | + print(f"An error occurred : {e}") |
| 121 | + config = self._get_secret(self.__get_secret_name_from_userdata()) |
| 122 | + self.configs = self.__add_defaults(config) |
| 123 | + log_level = 3 if self.configs['debug_mode'] else 1 |
| 124 | + self.__setup_vsockproxy(log_level) |
| 125 | + self.__run_config_server(log_level) |
| 126 | + self.__run_socks_proxy(log_level) |
| 127 | + |
| 128 | + def run_compute(self): |
| 129 | + self._setup_auxilaries() |
| 130 | + self._validate_auxilaries() |
| 131 | + command = [ |
| 132 | + "nitro-cli", "run-enclave", |
| 133 | + "--eif-path", "/opt/uid2operator/uid2operator.eif", |
| 134 | + "--memory", self.config['enclave_memory_mb'], |
| 135 | + "--cpu-count", self.config['enclave_cpu_count'], |
| 136 | + "--enclave-cid", 42, |
| 137 | + "--enclave-name", "uid2operator" |
| 138 | + ] |
| 139 | + if self.config['debug']: |
| 140 | + command+=["--debug-mode", "--attach-console"] |
| 141 | + subprocess.run(command, check=True) |
| 142 | + |
| 143 | + def cleanup(self): |
| 144 | + describe_output = subprocess.check_output(["nitro-cli", "describe-enclaves"], text=True) |
| 145 | + enclaves = json.loads(describe_output) |
| 146 | + enclave_id = enclaves[0].get("EnclaveID") if enclaves else None |
| 147 | + if enclave_id: |
| 148 | + subprocess.run(["nitro-cli", "terminate-enclave", "--enclave-id", enclave_id]) |
| 149 | + print(f"Enclave with ID {enclave_id} has been terminated.") |
| 150 | + else: |
| 151 | + print("No enclave found or EnclaveID is null.") |
| 152 | + |
| 153 | + def kill_process(self, process_name): |
| 154 | + try: |
| 155 | + result = subprocess.run( |
| 156 | + ["pgrep", "-f", process_name], |
| 157 | + stdout=subprocess.PIPE, |
| 158 | + text=True, |
| 159 | + check=False |
| 160 | + ) |
| 161 | + if result.stdout.strip(): |
| 162 | + for pid in result.stdout.strip().split("\n"): |
| 163 | + os.kill(int(pid), signal.SIGKILL) |
| 164 | + print(f"{process_name} exited") |
| 165 | + else: |
| 166 | + print(f"Process {process_name} not found") |
| 167 | + except Exception as e: |
| 168 | + print(f"Failed to shut down {process_name}: {e}") |
| 169 | + |
| 170 | +if __name__ == "__main__": |
| 171 | + parser = argparse.ArgumentParser() |
| 172 | + parser.add_argument("-o", "--operation", required=False) |
| 173 | + args = parser.parse_args() |
| 174 | + ec2 = EC2() |
| 175 | + if args.operation and args.operation == "stop": |
| 176 | + ec2.cleanup() |
| 177 | + [ec2.kill_process(process) for process in ["vsockpx", "sockd", "vsock-proxy", "nohup"]] |
| 178 | + else: |
| 179 | + ec2.run_compute() |
0 commit comments