163 lines
4.6 KiB
Python
163 lines
4.6 KiB
Python
|
"""
|
||
|
Python client based on https://pypi.org/project/websocket-client/
|
||
|
|
||
|
@author hornetfighter515
|
||
|
@author Todorov-Lachezar
|
||
|
"""
|
||
|
|
||
|
import websocket
|
||
|
import _thread
|
||
|
import time
|
||
|
import ssl
|
||
|
import yaml
|
||
|
import json
|
||
|
import rsa
|
||
|
import random # for message id generation
|
||
|
import cipher
|
||
|
|
||
|
ext = False
|
||
|
authed = False
|
||
|
clis = []
|
||
|
messages = {}
|
||
|
|
||
|
def authed_ext_send(ws, msg):
|
||
|
msg = rsa.encrypt(msg.encode(), clis[0]).decode('latin')
|
||
|
message_id = random.random()
|
||
|
if len(clis) <= 1: # if there are no IoTs
|
||
|
# send the message once-encrypted
|
||
|
message = cipher.gen_comm_msg(0, message_id, msg)
|
||
|
ws.send(message)
|
||
|
return
|
||
|
# otherwise, slice and twice encrypt it
|
||
|
twice_es = []
|
||
|
modifier = len(msg)/(len(clis)-1)
|
||
|
for i in range(1,len(clis)):
|
||
|
start = int( (i-1) * modifier)
|
||
|
end = int(i*modifier)
|
||
|
# makes a slice of the encrypted message, and indicates the intended
|
||
|
# recipient
|
||
|
contents = rsa.encrypt(msg[start:end].encode('latin'),
|
||
|
clis[i]).decode('latin')
|
||
|
message = cipher.gen_comm_msg(i, message_id, contents)
|
||
|
twice_es.append(message)
|
||
|
for m in twice_es:
|
||
|
ws.send(m)
|
||
|
|
||
|
|
||
|
def authed_int_send(ws, msg):
|
||
|
to = int(input(f"Who's it to? 0-{len(clis)}>"))
|
||
|
e_msg = rsa.encrypt(msg.encode(), clis[int(to)]).decode('latin')
|
||
|
message_id = random.random()
|
||
|
msg = cipher.gen_comm_msg(to, message_id, e_msg)
|
||
|
ws.send(msg)
|
||
|
|
||
|
|
||
|
def init_send(ws):
|
||
|
message = cipher.gen_pub_key_msg(pub['n'], pub['e'])
|
||
|
ws.send(message)
|
||
|
|
||
|
|
||
|
def send(ws, msg):
|
||
|
try:
|
||
|
if ext:
|
||
|
authed_ext_send(ws, msg)
|
||
|
else:
|
||
|
authed_int_send(ws, msg)
|
||
|
except Exception as e:
|
||
|
print(f'Sending message {msg} failed: {e}')
|
||
|
|
||
|
|
||
|
def on_message(ws, message):
|
||
|
msg = json.loads(message)
|
||
|
global authed
|
||
|
if msg['type'] == 'pub_key':
|
||
|
authed = True
|
||
|
if ext:
|
||
|
k = rsa.PublicKey(msg['available_cli']['n'],
|
||
|
msg['available_cli']['e'])
|
||
|
if len(clis) == 0:
|
||
|
clis.append(k)
|
||
|
else:
|
||
|
clis[0] = k
|
||
|
for i,p in msg['available_iots'].items():
|
||
|
k = rsa.PublicKey(p['n'], p['e'])
|
||
|
clis.append(k)
|
||
|
else:
|
||
|
k = rsa.PublicKey(msg['available_cli']['n'],
|
||
|
msg['available_cli']['e'])
|
||
|
clis.append(k)
|
||
|
|
||
|
elif msg['type'] == 'message':
|
||
|
# check if it's only a partial message
|
||
|
try:
|
||
|
if 'from' in msg:
|
||
|
if msg['id'] not in messages:
|
||
|
messages[msg['id']] = {}
|
||
|
messages[msg['id']][msg['from']] = msg['contents']
|
||
|
contents = None
|
||
|
for i in range(0, 3):
|
||
|
if i in messages[msg['id']]:
|
||
|
if contents is None:
|
||
|
contents = messages[msg['id']][i].encode('latin')
|
||
|
else:
|
||
|
contents += messages[msg['id']][i].encode('latin')
|
||
|
if len(contents) == 256:
|
||
|
m = rsa.decrypt(contents, priv).decode()
|
||
|
else:
|
||
|
m = ' '
|
||
|
else:
|
||
|
contents = msg['contents'].encode('latin')
|
||
|
m = rsa.decrypt(contents, priv).decode()
|
||
|
if m is not None:
|
||
|
print(f"\033[F\033[1G<< {m}\n> ")
|
||
|
except Exception as e:
|
||
|
print(f'Failed to receive message due to {e}')
|
||
|
|
||
|
|
||
|
def on_error(ws, error):
|
||
|
print(error)
|
||
|
|
||
|
|
||
|
def on_close(ws, close_status_code, close_msg):
|
||
|
print(f"### closed. reason: {close_msg} ###")
|
||
|
|
||
|
|
||
|
def on_open(ws):
|
||
|
def run(*args):
|
||
|
init_send(ws)
|
||
|
running = True
|
||
|
while running:
|
||
|
outbound = input(">")
|
||
|
if outbound == 'q':
|
||
|
running = False
|
||
|
else:
|
||
|
send(ws, outbound)
|
||
|
ws.close()
|
||
|
_thread.start_new_thread(run, ())
|
||
|
|
||
|
|
||
|
def open_socket(url, port):
|
||
|
ws = websocket.WebSocketApp(f"ws://{url}:{port}",
|
||
|
on_open=on_open,
|
||
|
on_message=on_message,
|
||
|
on_error=on_error,
|
||
|
on_close=on_close)
|
||
|
ws.run_forever(sslopt={"cert_reqs": ssl.CERT_NONE})
|
||
|
|
||
|
|
||
|
if __name__ == "__main__":
|
||
|
with open('src/config.yml', 'r') as f:
|
||
|
conf = yaml.safe_load(f)
|
||
|
|
||
|
e = input("Are you outside of the network? (Y/n)")
|
||
|
if e == 'y' or e == 'Y' or e == 'yes' or e == 'Yes':
|
||
|
ext = True
|
||
|
ws_conf = conf['ws']['ext']
|
||
|
else:
|
||
|
ws_conf = conf['ws']['int']
|
||
|
print('Generating keys...')
|
||
|
pub, priv = rsa.newkeys(2048)
|
||
|
print('Keys generated')
|
||
|
open_socket(ws_conf['url'], ws_conf['port'])
|
||
|
|