VsockStream class from vsock-parent.py

vsock-parent.py

class VsockStream:
    """Client"""
    def __init__(self, conn_tmo=15):
        self.conn_tmo = conn_tmo
        self.parent_private_key = None
        self.parent_public_key = None
        self.enclave_private_key = None
        

    def connect(self, endpoint):
        """Connect to the remote endpoint"""
        self.sock = socket.socket(socket.AF_VSOCK, socket.SOCK_STREAM)
        self.sock.settimeout(self.conn_tmo)
        self.sock.connect(endpoint)
        

    def send_keys_parent(self):
        print('Loading keys.')
        (self.parent_public_key, self.parent_private_key) = load_rsa_keys()
        length = pack('>Q', len(self.parent_public_key._save_pkcs1_pem()))
        self.sock.sendall(length)
        self.sock.sendall(self.parent_public_key._save_pkcs1_pem())
        print('Keys sent from parent')

    def send_image_parent(self, endpoint):
        encrypted_key = encrypt_image('basal_cell_carcinoma_example.png', 'enclave_public_key_received.pem')
        with open('basal_cell_carcinoma_example.png.encrypted', 'rb') as f:
            image_contents = f.read()

        length = pack('>Q', len(image_contents))
        print(f'Sending image of length {str(len(image_contents))}')
        while True:
            try:
                self.sock.sendall(length)
                print('Length message sent')
                self.sock.sendall(image_contents)
                break
            except socket.timeout:
                time.sleep(2)
    
        length = pack('>Q', len(encrypted_key))
        print('Sending symmetric key of length: ', str(len(encrypted_key)))
        self.connect(endpoint)
        self.sock.sendall(length)
        self.sock.sendall(encrypted_key)

        self.sock.close()
            

def client_handler(args):

    client = VsockStream()
    endpoint = (args.cid, args.port)
    client.connect(endpoint)

    client.send_keys_parent()
    
    client.connect(endpoint)
    client.send_image_parent(endpoint)