This post is part two of a multi-part series. If you haven’t already, please read part one here: Encrypted Chat: Part I
Welcome Back
After gaining an understanding of the concepts behind my encrypted chat server, we’re now ready to take a closer look at the code to see how the client and server operate.
Part Two: The Code
These are the files that make up the encrypted chat project:
File Description Category
----------------------------------------------------------------
client.py Chat Client Networking
server.py Chat Server Networking, Control
dhke.py Diffie-Hellman Key Exchange Crypto
cipher.py AES Encryption and Decryption Crypto
cli.py Command-Line Interface Interface
The client and the server share a lot of the same code, especially when it comes to the encryption and decryption of messages. However, since the server handles the majority of the processing, we’ll start there.
I’m only going to be showing the important bits so therefore some trivial code will be omitted (indicated by …). The full code can be found here.
2A. Server Code - Key Exchange
server.py
# Inside Server Class
def __init__(self, host='127.0.0.1', port=DEFAULT_PORT):
...
self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self.dh_params = M2DH.gen_params(DH_SIZE, 2)
...
We start by creating a new socket that the server will communicate on.
self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
We indicate the use of the IPv4 address family (e.g. 0x00sec.org, 74.125.136.94) with socket.AF_INET, and the use of TCP with socket.SOCK_STREAM.
Next, we generate the parameters for the Diffie-Hellman key exchange with the DH module from m2crypto.
self.dh_params = M2DH.gen_params(DH_SIZE, 2)
Remember that we needed to create two public values: a prime modulus p and prime base g. When we indicate the size (in bits) of the desired prime (a minimum of 2048 bits for the shared prime is recommended) and which generator we want to use, this function returns a class containing p and g. It does this using OpenSSL, a “robust, commercial-grade, and full-featured toolkit for the Transport Layer Security (TLS) and Secure Sockets Layer (SSL) protocols” and a “general-purpose cryptography library.” You’ll find many other programs use this library, such as OpenVPN.
We can now wait for new clients to connect, with which we will attempt Diffie-Hellman key exchange and listen for their incoming encrypted messages.
while True:
connection, address = self.socket.accept()
client = Client(self, connection, address)
...
self.clients.append(client)
...
threading.Thread(target=self.listen, args=(client, )).start()
We use our socket to accept the next incoming connection (this function is blocking) and initialize it as a new client. Let’s take a look at what happens when we do that:
# Inside Client Class (this is the client object the server uses)
def __init__(self, server, connection, address):
...
self.key = self.dh(server.dh_params)
The client class takes in a server (to which the client belongs), a socket connection, and address in the format (ip_address, port). It then calls dh() to generate the shared key.
# Inside Client Class
def dh(self, dh_params):
"""
Perform Diffie-Hellman Key Exchange with a client.
:param dh_params: p and g generated by DH
:return shared_key: shared encryption key for AES
"""
# p: shared prime
p = DH.b2i(dh_params.p)
# g: primitive root modulo
g = DH.b2i(dh_params.g)
# a: randomized private key
a = DH.gen_private_key()
# Generate public key from p, g, and a
public_key = DH.gen_public_key(g, a, p)
# Create a DH message to send to client as bytes
dh_message = bytes(DH(p, g, public_key))
self.connection.sendall(dh_message)
# Receive public key from client as bytes
try:
response = self.connection.recv(LEN_PK)
except ConnectionError:
print("Key Exchange with {} failed".format(self.address[0]))
return None
client_key = DH.b2i(response)
# Calculate shared key with newly received client key
shared_key = DH.get_shared_key(client_key, a, p)
return shared_key
First, the method converts the p and g parameters from bytes to integers with DH.b2i, located in dhke.py. It then generates a random private key a (also an integer). The server’s public key for this exchange is calculated by raising g to the power of a and modulo-ing the result by p with DH.gen_public_key().
def gen_public_key(g, private, p):
# g^private % p
return pow(g, private, p)
Now it can use the values p, g, and public_key to construct a public message to send to the client.
**Note: the client needs these three values to generate its own public key
The following line creates a new DH object and converts it to bytes:
dh_message = bytes(DH(p, g, public_key))
We can see how the three variables are encoded into bytes in the __bytes__ method of DH:
def __bytes__(self):
"""
Convert DH message to bytes.
:return: packaged DH message as bytes
+-------+-----------+------------+
| Prime | Generator | Public Key |
| 1024 | 16 | 1024 |
+-------+-----------+------------+
"""
prm = self.package(self.p, LEN_PRIME)
gen = self.package(self.g, LEN_GEN)
pbk = self.package(self.pk, LEN_PK)
return prm + gen + pbk
Since we need a standardized message format for the client to unpack, the first 1024 bytes belong to p, the following 16 to g, and the last 1024 to the public key. The package method just converts the integer variable to bytes and adds padding until it hits the desired length:
def package(i, length):
"""
Package an integer as a bytes object of length "length".
:param i: integer to be package
:param length: desired length of the bytes object
:return: bytes representation of the integer
"""
# Convert i to hex and remove '0x' from the left
i_hex = hex(i)[2:]
# Make the length of i_hex a multiple of 2
if len(i_hex) % 2 != 0:
i_hex = '0' + i_hex
# Convert hex string into bytes
i_bytes = binascii.unhexlify(i_hex)
# Check to make sure bytes to not exceed the max length
len_i = len(i_bytes)
if len_i > length:
raise InvalidDH("Length Exceeds Maximum of {}".format(length))
# Generate padding for the remaining space on the left
i_padding = bytes(length - len_i)
return i_padding + i_bytes
With the message all packaged up, we send it over to the client and wait for a response with the client’s public key:
self.connection.sendall(dh_message)
try:
response = self.connection.recv(LEN_PK)
except ConnectionError:
print("Key Exchange with {} failed".format(self.address[0]))
return None
client_key = DH.b2i(response)
We then convert the response from bytes to an integer and pass it into the DH.get_shared_key() method with our own private key a and the shared prime p:
shared_key = DH.get_shared_key(client_key, a, p)
get_shared_key() calculates (client_key ^ a) % p, converts the results to a hex string, and passes it through sha256 to standardize its length:
def get_shared_key(public, private, p):
"""
Calculate a shared key from a foreign public key, a local private
key, and a shared prime.
:param public: public key as an integer
:param private: private key as an integer
:param p: prime number
:return: shared key as a 256-bit bytes object
"""
s = pow(public, private, p)
s_hex = hex(s)[2:]
# Make the length of s_hex a multiple of 2
if len(s_hex) % 2 != 0:
s_hex = '0' + s_hex
# Convert hex to bytes
s_bytes = binascii.unhexlify(s_hex)
# Hash and return the hex result
return sha256(s_bytes).digest()
Finally, we’ve got a shared key! That’s pretty cool, but of course it was a little more complicated in practice than it was in theory. Now, where were we?
With the key declared, we have just finished initializing a client within the server.
client = Client(self, connection, address)
...
# Add client to list of clients on server
self.clients.append(client)
...
# Listen for incoming messages from client
threading.Thread(target=self.listen, args=(client, )).start()
In a new thread we listen for income encrypted messages from the client, decrypt them, and broadcast them (re-encrypted) to every other client on the server with self.listen.
This post is already long enough as-is, so I’ll stop here for now. Check in next time when we’ll look at how the server processes messages and maybe see what’s going on with the client. Thanks for reading!
