def stage_three(self, stage_three_msg_main): """ Gets the public_key_main_dh, calculates the shared key and saves it. sends a random ACK message that will be encrypted and signed with the shared key. :param stage_three_msg_main (Bytes): public_key_main_dh :return Ack message (Bytes): A random ack message :return shared_key (Str): The shared_key """ self.public_key_main_dh = int(stage_three_msg_main.decode()) self.shared_key = DH.gen_shared_key(self.private_key_edge_dh, self.public_key_main_dh) ack = ''.join(random.choice(string.printable) for i in range(128)).encode() signature = zRSA.sign_data(ack, self.private_rsa_edge) return ack, signature, self.shared_key
def stage_two(self, stage_two_msg_edge): """ Gets the public_key_edge_dh and saves it. Creates a DH key set for main and sends the public_key_main_dh. Calculates the shared key and returns it. :param stage_two_msg_edge (Bytes): public_key_edge_dh :return stage_two_msg_main (Bytes): public_key_main_dh :return shared_key: the final shared key """ self.public_key_edge_dh = int(stage_two_msg_edge.decode()) self.private_key_main_dh, self.public_key_main_dh = DH.gen_key_set() self.shared_key = DH.gen_shared_key(self.private_key_main_dh, self.public_key_edge_dh) stage_two_msg_main = str(self.public_key_main_dh).encode() return stage_two_msg_main, self.shared_key