cryptopals/TCPClient.py at main · IOKernel/cryptopals · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
import socket
import threading
from utils import power_mod, Random
from time import time
from hashing import sha1
from os import urandom
from aes import aes_cbc_decrypt, aes_cbc_encrypt
from padding import pkcs7_pad, pkcs7_unpad

class Client():
    def __init__(self, HOST = 'localhost', PORT = 10000):
        # create a TCP socket
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.host = HOST
        self.port = PORT
        # initializing needed parameters
        self._private_key = 0 # private key a|b
        self.public_key = 0 # public key A|B
        self.shared_key = b"" # share key from sha1(s)
        self.other = "" # other user's nickname

        # connect to the server
        self._connect()

        # Starting Threads For Listening And Writing
        receive_thread = threading.Thread(target=self.receive)
        receive_thread.start()

        write_thread = threading.Thread(target=self.write)
        write_thread.start()

    def initiate_dh(self, p, g, a):
        # setting DH parameters from user and sending them
        self.p = p
        self.g = g
        self._private_key = power_mod(a, 1, p)
        self.public_key = power_mod(g, self._private_key, p)
        init_str = f"[p, g, A] = {str(p)}, {str(g)}, {str(self.public_key)}".encode()
        print("\U00002705 Sending DH parameters...")
        self.sock.sendall(init_str)

    def _connect(self):
        # initiate server connection
        try:
            self.sock.connect((self.host, self.port))
            print(f"Connected to {self.host}:{self.port}")

        except:
            print('server is off')

    # Sending Messages To Server
    def write(self):
        while True:
            message = input('')
            try:
                if self.shared_key:
                # dh exchange is over, encrypt messages
                    iv = urandom(16)
                    padded = pkcs7_pad(message)
                    encrypted_msg = aes_cbc_encrypt(padded, self.shared_key, iv)
                    self.sock.send(encrypted_msg + iv)

                else:
                    self.sock.send(message.encode())

            except KeyboardInterrupt:
                print('closing connection.')
                exit()

    # Listening to Server and Sending Nickname
    def receive(self):
        while True:
            # Receive Message From Server
            message = self.sock.recv(1024)

            # initiation of dh and encryption
            if message[:11] ==  b"[p, g, A] =":
                print("\U00002705 Received DH parameters, sending public key...")
                msg = message.decode()
                self.p, self.g, A = [int(n.strip(',')) for n in msg.split()[4:]]
                b = Random(int(time())).random()
                self._private_key = power_mod(b, 1, self.p)
                self.public_key = power_mod(self.g, self._private_key, self.p)
                s = power_mod(A, self._private_key, self.p)
                self.shared_key = sha1(str(s).encode()).bytes()[0:16]
                print("\U00002705 Connection is now encrypted.")
                self.sock.sendall(f"B = {str(self.public_key)}".encode())

            elif message[:8] == b"[p, g] =":
                print("\U00002705 Received DH parameters, sending ACK...")
                msg = message.decode()
                self.p, self.g = [int(n.strip(',')) for n in msg.split()[3:]]
                b = Random(int(time())).random()
                self._private_key = power_mod(b, 1, self.p)
                self.public_key = power_mod(self.g, self._private_key, self.p)
                self.sock.send(b"ACK")

            # receive other's public key
            elif message[:3] == b"B =":
                print("\U00002705 Received public key!")
                B = int(message.split()[2])
                s = power_mod(B, self._private_key, self.p)
                self.shared_key = sha1(str(s).encode()).bytes()[0:16]
                print("\U00002705 Connection is now encrypted.")

            elif message[:3] == b"A =":
                print("\U00002705 Received public key!")
                A = int(message.split()[2])
                s = power_mod(A, self._private_key, self.p)
                self.shared_key = sha1(str(s).encode()).bytes()[0:16]
                print("\U00002705 Connection is now encrypted.")
                self.sock.sendall(f"B = {str(self.public_key)}".encode())

            # print join statement and get other's Nicknname
            elif message[2:6] == b"join":
                if not self.other:
                    if message.decode()[0] == 'A':
                        self.other = 'B'
                    elif message.decode()[0] == 'B':
                        self.other = 'A'
                print(message.decode())


            elif message[0:3] == b"ACK":
                print("\U00002705 Received ACK, sending A...")
                self.sock.sendall(f"A = {str(self.public_key)}".encode())

            # if connection closes
            elif message == b'':
                print("connection lost!")
                self.sock.close()
                break

            # After key exchange is done, messages are encrypted using the shared key
            # decrypt the message using the iv appended to the text and shared key
            else:
                print(f"{self.shared_key = }")
                iv = message[-16:]
                pt_padded = aes_cbc_decrypt(message[:-16], self.shared_key, iv)
                pt = pkcs7_unpad(pt_padded)
                print(f"{self.other}: {pt.decode()}")