generated from mwc/lab_subrosa
129 lines
3.7 KiB
Python
129 lines
3.7 KiB
Python
from pathlib import Path
|
|
from api import SubRosaAPI, APIError
|
|
from datetime import datetime
|
|
import sys
|
|
sys.path.insert(0, ".")
|
|
from encryption import PrivateKey, PublicKey
|
|
|
|
class SubRosaUI:
|
|
"""A user interface to the SubRosa chat system.
|
|
"""
|
|
def __init__(self, private_key_file, url):
|
|
self.api = SubRosaAPI(url)
|
|
if Path(private_key_file).exists():
|
|
self.private_key = PrivateKey.load(private_key_file)
|
|
else:
|
|
self.private_key = PrivateKey.generate()
|
|
self.private_key.save(private_key_file)
|
|
self.public_key = self.private_key.get_public_key()
|
|
|
|
def run(self, username):
|
|
"""Runs the client UI.
|
|
"""
|
|
self.username = username
|
|
self.user = self.get_or_create_user(username)
|
|
print("Welcome to SubRosa.")
|
|
while True:
|
|
choice = self.choose_action_from_menu()
|
|
if choice == 'VIEW':
|
|
self.show_messages()
|
|
elif choice == 'SEND':
|
|
self.send_message()
|
|
print("=" * 80)
|
|
|
|
def get_or_create_user(self, username):
|
|
"""Gets user info from the server, or creates the user.
|
|
If the user already exists, checks that our public key matches
|
|
the server's public key for the user.
|
|
"""
|
|
try:
|
|
user = self.api.get_user(username)
|
|
if user['public_key'] != str(self.public_key):
|
|
raise ValueError(f"Invalid key for {username}")
|
|
except APIError:
|
|
user = self.api.create_user(username, str(self.public_key))
|
|
return user
|
|
|
|
def choose_action_from_menu(self):
|
|
"""Gets a choice from the user.
|
|
"""
|
|
print("What would you like to do?")
|
|
print("1. See messages")
|
|
print("2. Send a message")
|
|
while True:
|
|
choice = input("> ").strip()
|
|
if choice == '1':
|
|
return 'VIEW'
|
|
elif choice == '2':
|
|
return 'SEND'
|
|
|
|
def show_messages(self):
|
|
messages = self.api.get_messages(self.username)
|
|
|
|
for i, message in enumerate(messages):
|
|
ciphertext = message["ciphertext"]
|
|
sender = message["sender"]
|
|
|
|
try:
|
|
plaintext = self.private_key.decrypt(ciphertext)
|
|
except Exception:
|
|
plaintext = "[Could not decrypt]"
|
|
|
|
print("-" * 80)
|
|
print(f"{i}. From {sender}:")
|
|
print(plaintext)
|
|
|
|
print("=" * 80)
|
|
|
|
def send_message(self):
|
|
recipient = input("Username of recipient: ")
|
|
message = input("Message: ")
|
|
|
|
# get recipient public key
|
|
recipient_public_key = self.api.get_public_key(recipient)
|
|
|
|
# encrypt message
|
|
ciphertext = recipient_public_key.encrypt(message)
|
|
|
|
# get current time
|
|
time_sent = get_current_time()
|
|
|
|
# sign the time
|
|
time_sent_signature = self.private_key.sign(time_sent)
|
|
|
|
# send message
|
|
self.api.send_message(
|
|
self.username,
|
|
recipient,
|
|
ciphertext,
|
|
time_sent,
|
|
time_sent_signature
|
|
)
|
|
|
|
print("-" * 80)
|
|
print("Message sent!")
|
|
print("=" * 80)
|
|
|
|
|
|
def get_current_time(self):
|
|
return datetime.utcnow().isoformat()
|
|
|
|
def main():
|
|
from argparse import ArgumentParser
|
|
parser = ArgumentParser()
|
|
parser.add_argument("username", help="Existing or new username.")
|
|
parser.add_argument("-k", "--key",
|
|
default="subrosa_private_key.pem",
|
|
help="Private key file. Will be created if it does not exist."
|
|
)
|
|
parser.add_argument("-s", "--server-url",
|
|
default="https://subrosa.makingwithcode.org",
|
|
help='Server URL. Use "http://127.0.0.1:8000" for a local server.'
|
|
)
|
|
args = parser.parse_args()
|
|
ui = SubRosaUI(args.key, args.server_url)
|
|
ui.run(args.username)
|
|
|
|
if __name__ == '__main__':
|
|
main()
|