generated from mwc/lab_subrosa
107 lines
3.6 KiB
Python
107 lines
3.6 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):
|
|
"""Fetches all messages and displays them.
|
|
Note: The messages arrive from the server encrypted.
|
|
"""
|
|
messages = self.api.get_messages(self.username)['messages']
|
|
if messages:
|
|
for i, message in enumerate(messages):
|
|
print('-' * 80)
|
|
print(f"{i}. From {message['sender']}:")
|
|
print(message['ciphertext'])
|
|
else:
|
|
print(f"No messages for {self.username}.")
|
|
|
|
def send_message(self):
|
|
"""Sends a message.
|
|
"""
|
|
try:
|
|
recipient_name = input("Username of recipient: ").strip()
|
|
recipient = self.api.get_user(recipient_name)
|
|
except APIError:
|
|
print(f"No user named {recipient_name}.")
|
|
return
|
|
plaintext = input("Message: ")
|
|
print('-' * 80)
|
|
print(f"Sorry, couldn't send a message to {recipient_name}; this method isn't implemented yet.")
|
|
|
|
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()
|