def GetAttestation(self, request, context):
"""
Calls get_remote_report_with_public_key()
"""
# Get a reference to the existing enclave
enclave_reference = xgb.Enclave(create_enclave=False)
# Get report from enclave
enclave_reference.get_remote_report_with_pubkey()
pem_key, key_size, remote_report, remote_report_size = enclave_reference.get_report_attrs(
)
return remote_attestation_pb2.Report(
pem_key=pem_key,
key_size=key_size,
remote_report=remote_report,
remote_report_size=remote_report_size)
def run(channel_addrs, key_path, keypair):
"""
The client will make 4 calls to the server that will run computation
1. A call to retrieve the attestation report from the server. The client will use this report
to verify that the it can trust the server.
2. A call to send the symmetric key used to encrypt the data to the server.
3. A call to commence computation.
"""
# Perform attestation and send the client key to each node in the enclave cluster
for channel_addr in channel_addrs:
print("Connecting to", channel_addr)
# Get remote report from enclave
with grpc.insecure_channel(channel_addr) as channel:
stub = remote_attestation_pb2_grpc.RemoteAttestationStub(channel)
response = stub.GetAttestation(
remote_attestation_pb2.Status(status=1))
pem_key = response.pem_key
key_size = response.key_size
remote_report = response.remote_report
remote_report_size = response.remote_report_size
print("Report received from remote enclave")
# Verify report
enclave_reference = xgb.Enclave(create_enclave=False)
enclave_reference.set_report_attrs(pem_key, key_size, remote_report,
remote_report_size)
enclave_reference.verify_remote_report_and_set_pubkey()
print("Report successfully verified")
# Encrypt and sign symmetric key used to encrypt data
key_file = open(key_path, 'rb')
sym_key = key_file.read() # The key will be type bytes
key_file.close()
crypto_utils = xgb.CryptoUtils()
# Encrypt symmetric key
enc_sym_key, enc_sym_key_size = crypto_utils.encrypt_data_with_pk(
sym_key, len(sym_key), pem_key, key_size)
print("Encrypted symmetric key")
# Sign encrypted symmetric key
sig, sig_len = crypto_utils.sign_data(keypair, enc_sym_key,
enc_sym_key_size)
print("Signed ciphertext")
# Send data key to the server
with grpc.insecure_channel(channel_addr) as channel:
stub = remote_attestation_pb2_grpc.RemoteAttestationStub(channel)
print("Sending key...")
response = stub.SendKey(
remote_attestation_pb2.DataMetadata(enc_sym_key=enc_sym_key,
key_size=enc_sym_key_size,
signature=sig,
sig_len=sig_len))
print("Symmetric key for data sent to server")
print("Waiting for training to finish...")
# Open up a channel to each node to signal start
channels = []
for channel_addr in channel_addrs:
# Signal start
channels.append(grpc.insecure_channel(channel_addr))
# Store futures in a list
# Futures hold the result of asynchronous calls to each gRPC server
futures = []
for channel in channels:
stub = remote_attestation_pb2_grpc.RemoteAttestationStub(channel)
# Asynchronous calls to start job on each node
response_future = stub.SignalStartCluster.future(
remote_attestation_pb2.ClusterParams(num_workers=2))
futures.append(response_future)
results = []
for future in futures:
results.append(future.result().status)
# If any node returned a non zero exit status
if sum(results) == 0:
print("Training succeeded! Encrypted model has been saved.")
else:
print("Training failed")
username = "******"
HOME_DIR = os.path.dirname(os.path.realpath(__file__)) + "/../../"
sym_key_file = HOME_DIR + "demo/data/key_zeros.txt"
pub_key_file = HOME_DIR + "demo/data/userkeys/private_user_1.pem"
cert_file = HOME_DIR + "demo/data/usercrts/{0}.crt".format(username)
temp_name = HOME_DIR + "demo/data/temp_file.txt"
temp_enc_name = HOME_DIR + "demo/data/temp_file.txt.enc"
print("Init user parameters")
xgb.init_user(username, sym_key_file, pub_key_file, cert_file)
print("Creating enclave")
enclave = xgb.Enclave(HOME_DIR + "build/enclave/xgboost_enclave.signed")
# Remote Attestation
print("Remote attestation")
# Note: Simulation mode does not support attestation
# pass in `verify=False` to attest()
enclave.attest(verify=False)
print("Send private key to enclave")
enclave.add_key()
class TestUpdaters(unittest.TestCase):
@pytest.mark.skipif(**tm.no_sklearn())
def test_histmaker(self):
def run(channel_addr, key_path, keypair):
"""
The client will make 4 calls to the server that will run computation
1. A call to retrieve the attestation report from the server. The client will use this report
to verify that the it can trust the server.
2. A call to send the symmetric key used to encrypt the data to the server.
3. A call to commence computation.
"""
# Get remote report from enclave
with grpc.insecure_channel(channel_addr) as channel:
stub = remote_attestation_pb2_grpc.RemoteAttestationStub(channel)
response = stub.GetAttestation(remote_attestation_pb2.Status(status=1))
pem_key = response.pem_key
key_size = response.key_size
remote_report = response.remote_report
remote_report_size = response.remote_report_size
print("Report received from remote enclave")
# Verify report
enclave_reference = xgb.Enclave(create_enclave=False)
enclave_reference.set_report_attrs(pem_key, key_size, remote_report,
remote_report_size)
enclave_reference.verify_remote_report_and_set_pubkey()
print("Report successfully verified")
# Encrypt and sign symmetric key used to encrypt data
key_file = open(key_path, 'rb')
sym_key = key_file.read() # The key will be type bytes
key_file.close()
crypto_utils = xgb.CryptoUtils()
# Encrypt symmetric key
enc_sym_key, enc_sym_key_size = crypto_utils.encrypt_data_with_pk(
sym_key, len(sym_key), pem_key, key_size)
print("Encrypted symmetric key")
# Sign encrypted symmetric key
sig, sig_len = crypto_utils.sign_data(keypair, enc_sym_key,
enc_sym_key_size)
print("Signed ciphertext")
# Send data key to the server
with grpc.insecure_channel(channel_addr) as channel:
stub = remote_attestation_pb2_grpc.RemoteAttestationStub(channel)
response = stub.SendKey(
remote_attestation_pb2.DataMetadata(enc_sym_key=enc_sym_key,
key_size=enc_sym_key_size,
signature=sig,
sig_len=sig_len))
print("Symmetric key for data sent to server")
# Signal start
with grpc.insecure_channel(channel_addr) as channel:
stub = remote_attestation_pb2_grpc.RemoteAttestationStub(channel)
print("Waiting for training to finish...")
response = stub.SignalStart(remote_attestation_pb2.Status(status=1))
if response.status == 1:
print("Training succeeded! Decrypting predictions...")
enc_preds_serialized = response.predictions
num_preds = response.num_preds
enc_preds = proto_to_pointer(enc_preds_serialized)
preds = crypto_utils.decrypt_predictions(sym_key, enc_preds,
num_preds)
print("Predictions: ", preds)
else:
print("Training failed")
OE_ENCLAVE_FLAG_DEBUG = 1
OE_ENCLAVE_FLAG_SIMULATE = 2
print("Creating enclave")
HOME_DIR = os.getcwd() + "/../../../"
flags = OE_ENCLAVE_FLAG_RELEASE
# Uncomment below for enclave debug mode
flags |= OE_ENCLAVE_FLAG_DEBUG
# Uncomment below for enclave simulation mode
# flags |= OE_ENCLAVE_FLAG_SIMULATE
enclave = xgb.Enclave(HOME_DIR + "build/enclave/xgboost_enclave.signed",
flags=(flags))
# Remote Attestation
# print("Remote attestation")
# enclave.get_remote_report_with_pubkey()
# enclave.verify_remote_report_and_set_pubkey()
rabit_args = {
"DMLC_NUM_WORKER": os.environ.get("DMLC_NUM_WORKER"),
"DMLC_NUM_SERVER": os.environ.get("DMLC_NUM_SERVER"),
"DMLC_TRACKER_URI": os.environ.get("DMLC_TRACKER_URI"),
"DMLC_TRACKER_PORT": os.environ.get("DMLC_TRACKER_PORT"),
"DMLC_ROLE": os.environ.get("DMLC_ROLE"),
"DMLC_NODE_HOST": os.environ.get("DMLC_NODE_HOST")
}
import securexgboost as xgb
import os
DIR = os.path.dirname(os.path.realpath(__file__))
HOME_DIR = DIR + "/../../../"
print("Creating enclave")
enclave = xgb.Enclave(HOME_DIR + "build/enclave/xgboost_enclave.signed",
log_verbosity=0)
# Remote Attestation
print("Remote attestation")
# Note: Simulation mode does not support attestation
# pass in `verify=False` to attest()
enclave.attest()
username1 = "user1"
train_enc_1 = HOME_DIR + "demo/python/multiclient-distributed/data/u1_train.enc"
KEY_FILE_1 = HOME_DIR + "demo/python/multiclient-distributed/data/key1.txt"
PRIVATE_KEY_FILE_1 = HOME_DIR + "config/user1.pem"
CERT_FILE_1 = HOME_DIR + "config/user1.crt"
xgb.init_user(username1, KEY_FILE_1, PRIVATE_KEY_FILE_1, CERT_FILE_1)
print("Send private key to enclave")
enclave.add_key()
username2 = "user2"
train_enc_2 = HOME_DIR + "demo/python/multiclient-distributed/data/u2_train.enc"
test_enc_2 = HOME_DIR + "demo/python/multiclient-distributed/data/u2_test.enc"
def run(channel_addr, sym_key_file, priv_key_file, cert_file):
xgb.init_user(username, sym_key_file, priv_key_file, cert_file)
# Remote attestation
print("Remote attestation")
enclave_reference = xgb.Enclave(addr=channel_addr)
# Note: Simulation mode does not support attestation
# pass in `verify=False` to attest()
enclave_reference.attest()
print("Report successfully verified")
print("Send private key to enclave")
enclave_reference.add_key()
print("Creating training matrix")
dtrain = xgb.DMatrix({username: HOME_DIR + "demo/python/remote-control/data/train.enc"})
if not dtrain:
print("Error creating dtrain")
return
print("dtrain: " + dtrain.handle.value.decode("utf-8"))
print("Creating test matrix")
dtest = xgb.DMatrix({username: HOME_DIR + "demo/python/remote-control/data/test.enc"})
if not dtest:
print("Error creating dtest")
return
print("dtest: " + dtest.handle.value.decode("utf-8"))
print("Beginning Training")
# Set training parameters
params = {
"tree_method": "hist",
"n_gpus": "0",
"objective": "binary:logistic",
"min_child_weight": "1",
"gamma": "0.1",
"max_depth": "3",
"verbosity": "0"
}
# Train and evaluate
num_rounds = 5
print("Training...")
booster = xgb.train(params, dtrain, num_rounds)
print("booster: " + booster.handle.value.decode("utf-8"))
booster.save_model(HOME_DIR + "demo/python/remote-control/client/modelfile.model")
# Get encrypted predictions
print("\nModel Predictions: ")
predictions, num_preds = booster.predict(dtest, decrypt=False)
# Decrypt predictions
print(booster.decrypt_predictions(predictions, num_preds))
# Get fscores of model
print("\nModel Feature Importance: ")
print(booster.get_fscore())
