dist-iot-net/src/cli.py
hornet a85b08f050
Feature/py iot (#4)
* Changing files but having issues with space tabs.

* Removed config. Fixed indents.

* Removed config.

* Fixed test config. Fixed syntax in cli.

* Editing cli and hub to finally commit transactions.

* Added better logging logic.
Will clean up in the morning.

* Hub is more functional.

* Edited IoT to gen keys.

* Public keys are now generated, sent, and stored.
However, they are not transmitted when a new client joins.

* Added crypto testing

* Communication is now functional.
2021-11-28 19:00:46 +00:00

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'])