Implementing-cryptography-using-python

266 Chapter 9 ■ Mastering Cryptography Using Python
Figure 9.8: Final version
Modifying the Helper File
Now that you understand how the Diffie-Hellman protocol works, it is time to
implement the processes into the helper file. Examine the helper file that you
have created and implement the following code:
# encryption method used by all calls
def encrypt(message, usePKI, useDH, dhSecret):
if usePKI == True:
message = encrypt_rsa(message)
if useDH == True:
message = encrypt_dh(message, dhSecret)
return message
# decryption method used by all calls
def decrypt(message, usePKI, useDH, dhSecret):
if useDH == True:
message = decrypt_dh(message, dhSecret)
if usePKI == True:
message = decrypt_rsa(message)
return message
# delete all generated DH certs
def remove_dh_certs():
try:
os.remove("client_private_dh_key.pem")
os.remove("client_public_dh_key.pem")
os.remove("server_private_dh_key.pem")
os.remove("server_public_dh_key.pem")