-
Notifications
You must be signed in to change notification settings - Fork 0
/
peer.py
executable file
·405 lines (332 loc) · 11.2 KB
/
peer.py
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
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
#!/usr/bin/env python
"""
This module contains the executable for the p2p part of the chat
"""
import io
import ipaddress
import sys
from select import select
from socket import AF_INET, SOCK_STREAM, create_connection, socket
from sys import stdin
import capnp
import msgs_capnp
EXIT_OK = 0
EXIT_ERR = 1
incoming_port: int = 0
counter: int = 0
class PeerId():
"""
A socket ID object that is hashable
"""
inner: bytearray
def __init__(self, host: bytes, port: bytes):
self.inner = bytearray()
self.inner += host
self.inner += port
def __hash__(self):
num = int.from_bytes(self.inner, byteorder='big', signed=False)
return num
def __repr__(self):
return f"{self.__hash__()}"
class Peer:
"""
A peer object with its corresponding socket connection
"""
in_sock: socket
out_sock: socket
peer_id: PeerId
def __init__(
self,
peer_id: PeerId,
in_sock: socket = None,
out_sock: socket = None
):
"""
Create a peer given the peer_id
"""
self.peer_id = peer_id
if in_sock is not None:
self.in_sock = in_sock
if out_sock is not None:
self.out_sock = out_sock
def __hash__(self):
return self.peer_id.__hash__()
def __repr__(self):
return f"Peer(id: {self.peer_id.__hash__()})"
def send_msg(self):
pass
class Server:
"""
A server object connection
"""
sock: socket
buff: bytearray
def __init__(self, addr: str, port: str):
self.buff = bytearray(0)
try:
self.sock = create_connection((addr, int(port)))
except ValueError as e:
print(f"[x] Error while trying to parse arguments: {e}")
sys.exit(EXIT_ERR)
except OSError as e:
print(f"[x] Error while trying to connect to server: {e.strerror}")
sys.exit(EXIT_ERR)
def send_port(self):
"""
Once connected to the server, advertise which port this peer is
listening on
"""
global incoming_port
port_msg = msgs_capnp.PeerListeningPort.new_message()
port_msg.port = incoming_port
self.sock.send(port_msg.to_bytes())
def get_peers(self) -> dict[int, Peer]:
"""
Once connected to the server, get the connected peers and connect to
them
"""
# The first message received is a 4 bytes tag containing the size of
# the message
msg = self.sock.recv(4)
if len(msg) == 0:
print("[x] Connection unexpectedly died")
exit(EXIT_ERR)
# Count and size are declared here for clarity
#
# The size is the size received
# The count is the bytes received afterwards, used to check if the
# received message is complete
count = 0
size = 0
try:
size = int.from_bytes(msg, byteorder='big', signed=False)
except ValueError as e:
print(f"[x] Error while reading size, corrupted message {e}")
while count != size:
msg = self.sock.recv(4096)
count += len(msg)
if len(msg) == 0:
break
else:
self.buff += msg
# If we received less bytes than the server said it was going to send,
# we assume that the connection was aborted and terminate
if count != size:
print("[x] Connection unexpectedly died")
exit(EXIT_ERR)
sockets: dict[int, Peer] = {}
# Deserialize the message and connect to the peers sent by the server
addrs = msgs_capnp.ServerRpcMsg.from_bytes(self.buff)
# Craft port message to advertise to peers
port_msg = msgs_capnp.PeerListeningPort.new_message()
port_msg.port = incoming_port
wire_msg = port_msg.to_bytes()
for addr in addrs.addrs:
host = ipaddress.IPv4Address(addr.ip).exploded
port = addr.port
peer_id_host = addr.ip.to_bytes(4, byteorder='big', signed=False)
peer_id_port = addr.port.to_bytes(2, byteorder='big', signed=False)
peer_id = PeerId(peer_id_host, peer_id_port)
print(f"[i] Connecting to {host}:{port}")
try:
conn = create_connection((host, port))
conn.send(wire_msg)
sockets[conn.fileno()] = Peer(peer_id, out_sock=conn)
except OSError as e:
print(
f"[-] Error: Cannot connect to {host}:{port}, {e.strerror}"
)
return sockets
def handle_conn(
sock: socket,
select_map: dict[int, socket]
):
"""
Handles a new connection, registering its socket and creating its
corresponding outbound connection if it doesn't exist
"""
global incoming_port
global counter
msg = sock.recv(4096)
if msg == "":
print("[x] Error while reading from new connection")
return
select_map[sock.fileno()] = sock
def handle_msg(
sock: socket,
select_map: dict[int, socket]
):
"""
Handles peer messages, either unregistering peers or printing its content
"""
# Read the first 4096 bytes that tell us the size of the message
msg_size = sock.recv(4096)
if len(msg_size) == 0:
host, port = sock.getpeername()
del select_map[sock.fileno()]
print(f"[i] - {host}:{port} Disconnected")
return
try:
length = int.from_bytes(msg_size, byteorder='big', signed=False)
except ValueError as e:
print(f"[x] Error {e} while reading msg size")
return
# Initialize received to 0 to keep count of the bytes received
received = 0
# Initialize the buffer where we are going to save the message
msg = bytearray()
# Set timeout on socket just in case
sock.settimeout(0.75)
# Get the message or abort if nothing is received
while received < length:
new = sock.recv(4096)
if len(new) == 0:
host, port = sock.getpeername()
del select_map[sock.fileno()]
print(f"[i] - {host}:{port} Disconnected")
return
else:
msg += new
received += len(new)
# Clear the timeout
sock.settimeout(0)
# Deserialize msg
struct_msg = msgs_capnp.PeerMsg.from_bytes(msg)
# Handle text message
if struct_msg.type == "text":
print(struct_msg.content.text)
# Handle file message
elif struct_msg.type == "file":
# Destruct message fields
filename = struct_msg.content.file.filename
file_content = struct_msg.content.file.content
# Ask for permission to store file
print(
f"[?] File \"{filename}\" received," +
" do you want to save it? (y/n): ",
end=''
)
answer = input()
if answer == "y":
# Ask where to store the file
print("[?] What name do you want to give it?: ", end='')
filepath = input()
# Open and store file
try:
f = open(filepath, 'wb')
f.write(file_content)
except OSError as e:
print(f"[x] Error while writing to {e.strerror}")
return
print("[i] File written correctly")
def handle_stdin(
s: io.TextIOWrapper,
select_map: dict[int, socket],
listening: int
):
"""
Handles stdin input and dispatches it to its corresponding function
depending on the command
"""
msg = s.readline()
trimmed = msg.strip()
command, args = input_parse(trimmed)
# Text command is a simple message
if command == "text":
# Create the message to be sent
wire_msg = msgs_capnp.PeerMsg.new_message()
wire_msg.type = "text"
wire_msg.content.text = args
# Serialize message
bytes_msg = wire_msg.to_bytes()
size = len(bytes_msg)
# Broadcast the message
for p in select_map:
if p != 0 and p != listening:
sock = select_map[p]
sock.send(size.to_bytes(4096, byteorder='big', signed=False))
sock.send(bytes_msg)
# File command is to send files
elif command == "file":
# Open file path
try:
f = open(args, 'rb')
except OSError as e:
print(f"[x] Error while opening file \"{e.strerror}\"")
return
# Read file content into variable
file_content = f.read()
# Craft message to be sent
wire_msg = msgs_capnp.PeerMsg.new_message()
wire_msg.type = "file"
wire_msg.content.init('file')
wire_msg.content.file.filename = args
wire_msg.content.file.content = file_content
# Serialize message
bytes_msg = wire_msg.to_bytes()
size = len(bytes_msg)
# Broadcast the message
for p in select_map:
if p != 0 and p != listening:
sock = select_map[p]
sock.send(size.to_bytes(4096, byteorder='big', signed=False))
sock.send(bytes_msg)
print("[i] File sent correctly")
def input_parse(text: str) -> tuple[str, str]:
"""
Parses an input string into a command and a list of the arguments that it
takes
"""
ret: tuple[str, str] = ("", "")
if text[0] == '/':
# If the string starts with a / it is a command
com, _, rest = text.partition(' ')
ret = (com.removeprefix('/'), rest)
elif text[0] == '\\' and text[1] == '/':
# If the string starts with a \ and it is followed by / it is a message
# with an escaped /
ret = ('text', text.removeprefix('\\'))
else:
# Else it is just a message
ret = ('text', text)
return ret
def main():
global incoming_port
if len(sys.argv) != 3:
print(f"Usage:\n\t{sys.argv[0]} <host> <port>")
exit(1)
# Bind to random socket and listen to incoming messages
sock = socket(AF_INET, SOCK_STREAM)
sock.bind(("", 0))
sock.listen(8)
sock.setblocking(False)
incoming_port = sock.getsockname()[1]
# Connect to server
server = Server(sys.argv[1], sys.argv[2])
# Advertise port
server.send_port()
# Get peers
out_peers = server.get_peers()
# A hashmap from socketid to socket
select_map = {}
select_map[sock.fileno()] = sock
select_map |= map(lambda k: (k[0], k[1].out_sock), out_peers.items())
# Register stdin
select_map[stdin.fileno()] = stdin
while True:
selected, x, y = select(select_map, [], [])
for s in selected:
# This means that a new client connected
if s == sock.fileno():
conn, (host, port) = sock.accept()
print(f"[i] + {host}:{port} Connected")
handle_conn(conn, select_map)
# This means that the user input a message
elif s == stdin.fileno():
handle_stdin(stdin, select_map, sock.fileno())
# This means that a peer message was received
else:
f = select_map[s]
handle_msg(f, select_map)
if __name__ == "__main__":
main()