[CS 6740 Team Paracha and Sun] Updated client.py to fix the Linux character encoding incompatibility problem.

client.py

#!/usr/bin/env python
__author__ = "Team Talha and Shuwen"

import os
import sys
import socket
import getpass
import hashlib
import argparse
import configparser

from threading import Thread

import zmq

from lib.utils import (
    DEBUG,
    json_pack,
    json_unpack,
    nounce_generator,
    solve_challenge,
    hash_message
)
from lib.crypto import (
    decrypt_symmetric_key,
    deserialize_DH_public_key,
    DH_keys,
    encrypt_symmetric_key,
    encrypt_with_public_key,
    generate_DH_shared_key,
    serialize_DH_public_key,
    verify_with_public_key,
)


def prompt():
    """prompt for get commands from user"""
    sys.stdout.write("+> ")
    sys.stdout.flush()


def read_server_conf(conf):
    """Read conf file for server.

    The function read the config file input and return server side info or
    files.
    """
    config = configparser.ConfigParser()
    config.read(conf)

    if DEBUG:
        print(
            config["DEFAULT"]["IP"],
            config["DEFAULT"]["Port"],
            config["DEFAULT"]["ServerPK"],
        )

    return (
        config["DEFAULT"]["IP"],
        config["DEFAULT"]["Port"],
        config["DEFAULT"]["ServerPK"],
    )


def arg_parser():
    """Client side argument parsing.

    This is now a stripped version and we only have one arg -- server config
    file.
    """

    # parser

    parser = argparse.ArgumentParser(
        description="Client of Jethro's basic chat application.")

    # server public key
    parser.add_argument(
        "-c", action="store", required=True, help="config file", dest="conf")

    args = parser.parse_args()
    return args.conf


def user_login(
        server_ip,
        server_port,
        server_pk,
        contrib,
        secret,
        user_ip,
        user_port,
        context=None,
):
    """User login/authenticate handler.

    This function handles user authentication (establishing session, password
    etc), if anything fails the function will immediately exit the program.
    """
    if DEBUG:
        print("*** User Login entry point... ***\n")

    # get user name
    try:
        user_name = input("Input your user name: ")
        user_name = user_name.strip().lower()
    except Exception as err:
        print(
            "Unvalid input or other issue crashed the client: {}".format(err))
        sys.exit(1)

    # get user password
    try:
        passwd = getpass.getpass(prompt="Input your password: ")
    except Exception as err:
        print(
            "Unvalid input or other issue crashed the client: {}".format(err))
        sys.exit(1)

    # create context
    try:
        ctx = context or zmq.Context.instance()
    except Exception as err:
        print(
            "Trying to setup the zmq context failed. Exiting: {}".format(err))
        sys.exit(1)

    # setting up messgae queue
    try:
        socket = ctx.socket(zmq.REQ)
        socket.connect("tcp://" + str(server_ip) + ":" + str(server_port))
    except Exception as err:
        print("Bootstrapping Message Queue failed: {}".format(err))
        sys.exit(1)

    if DEBUG:
        print("tcp://" + server_ip + ":" + str(server_port))

    # CRAFT FIRST MSG: tell server user wants to login
    # ---------------------------------------------------------
    first_msg = json_pack(["msg"], ["Connect"])

    # Send it
    try:
        socket.send_json(first_msg)
    except Exception as err:
        print("socket error, sending the first message failed. Exiting: {}".
              format(err))
        sys.exit(1)

    # RECV SECOND MSG: server replies w/ a challenge
    # ---------------------------------------------------------
    try:
        reply_msg = socket.recv_json()
    except Exception as err:
        print("socket error, recving the fourth message failed. Exiting: {}".
              format(err))
        sys.exit(1)

    if DEBUG:
        print("*** RECVed SECOND msg... ***\n")
        print(reply_msg)

    reply_msg = json_unpack(reply_msg, ["challenge", "difficulty"])

    # CRAFT THIRD MSG: Respond with information needed to start a session
    # ----------------------------------------------------------------------------------------

    # Solve the challenge first
    try:
        challenge = reply_msg["challenge"]
        difficulty = int(reply_msg["difficulty"])
        challenge_response = solve_challenge(challenge, difficulty)
    except Exception as err:
        print("Challenge response failed. Exiting: {}".format(err))
        sys.exit(1)

    # Generate nounce N_1
    try:
        nonce_1 = nounce_generator()
    except Exception as err:
        print("Generate nounce failed: {}".format(err))
        sys.exit(1)
    try:
        contrib = serialize_DH_public_key(contrib).decode()

        json_data = json_pack(
            ["user", "contrib", "nounce", "user_ip", "user_port"],
            [user_name, contrib, nonce_1, user_ip, user_port],
        )

        # Integrity protected encryption done here
        ciphertext, nonce, key = encrypt_with_public_key(json_data, server_pk)
        encrypted_data = json_pack(["ciphertext", "nonce", "key"],
                                   [ciphertext, nonce, key])

        first_msg = json_pack(
            ["msg", "challenge", "ciphertext", "challenge_response"],
            [
                "ConnectWChallengeResponse",
                challenge,
                encrypted_data,
                challenge_response,
            ],
        )
    except Exception as err:
        print("Crafting the first message failed. Exiting: {}".format(err))
        sys.exit(1)

    if DEBUG:
        print("*** SENDing first msg... ***\n")
        print(first_msg)

    # SEND THIRD MSG: client tries to connect to server
    # ------------------------------------------------
    try:
        socket.send_json(first_msg)
    except Exception as err:
        print("socket error, sending the first message failed. Exiting: {}".
              format(err))
        sys.exit(1)

    # RECV FOURTH MSG: server replies w/ server's contribution
    # --------------------------------------------------------
    try:
        contribution_message_signed = socket.recv_json()
    except Exception as err:
        print("socket error, recving the fourth message failed. Exiting: {}".
              format(err))
        sys.exit(1)

    if DEBUG:
        print("*** RECVed fourth msg... ***\n")
        print(contribution_message_signed)

    # Different formats for the reply message
    try:
        contribution_message_signed = json_unpack(contribution_message_signed,
                                                  ["message", "signature"])
    except KeyError:
        contribution_message_signed = json_unpack(contribution_message_signed,
                                                  ["message"])

    # Verify the connection succeeded
    try:
        assert contribution_message_signed["message"] != "failure"
    except AssertionError as err:
        print(
            "\nFailed to establish a secure channel with server, probably because the user was already logged-in: {}"
            .format(err))
        sys.exit(1)

    # Verify the message actually came the server
    try:
        verify_with_public_key(
            contribution_message_signed["message"],
            contribution_message_signed["signature"],
            server_pk,
        )
    except Exception as err:
        print("Verifying the cipher failed. Exiting: {}".format(err))

    # Parse the message
    json_data = json_unpack(
        contribution_message_signed["message"],
        ["srv_contrib", "nounce_plus_one", "nounce_two"],
    )

    try:
        srv_contrib = deserialize_DH_public_key(
            json_data["srv_contrib"].encode())
        nounce_plus_one = json_data["nounce_plus_one"]
        nonce_two = json_data["nounce_two"]
    except Exception as err:
        print("Extracting info from the fourth message failed: {}".format(err))
        sys.exit(1)

    # Verify the freshness
    try:
        assert nounce_plus_one == nonce_1 + 1
    except Exception as err:
        print("\nVerifying the freshness failed: {}".format(err))
        sys.exit(1)

    # SESSION KEY ESTABLISHED
    # -----------------------
    # We now have the session key between client and server
    try:
        session_key = generate_DH_shared_key(srv_contrib, secret)
    except Exception as err:
        print(
            "\nCreating the seesion key between client and server failed: {}".
            format(err))
        sys.exit(1)

    if DEBUG:
        print("*** session key... ***\n")
        print(session_key)

    # FIFTH MSG: user login w/ password
    # ---------------------------------
    json_data = json_pack(
        ["user", "password", "nounce_two_plus_one"],
        [user_name, passwd, nonce_two + 1],
    )

    # Integrity protected encryption done here
    try:
        ciphertext, nonce = encrypt_symmetric_key(json_data, session_key)
    except Exception as err:
        print("\nEncrypt the symmetric key failed: {}".format(err))
        sys.exit(1)

    encrypted_data = json_pack(["ciphertext", "nonce"], [ciphertext, nonce])

    fifth_msg = json_pack(["msg", "user", "ciphertext"],
                          ["Login", user_name, encrypted_data])

    try:
        socket.send_json(fifth_msg)
    except Exception as err:
        print("Socket error, sending the fifth message failed: {}".format(err))
        sys.exit(1)

    if DEBUG:
        print("\n*** fifth msg data... ***\n")
        print(encrypted_data)
        print("\n****************\n")
        print(fifth_msg)

    # SIXTH MSG:
    # -------------------
    try:
        sixth_msg = socket.recv_json()
    except Exception as err:
        print(
            "\nsocket error, recving the sixth message failed: {}".format(err))
        sys.exit(1)

    # Decrypt incoming request using the session key
    encrypted_ciphertext = json_unpack(sixth_msg, ["ciphertext", "msg"])
    encrypted_ciphertext = json_unpack(encrypted_ciphertext["ciphertext"],
                                       ["ciphertext", "nonce"])
    try:
        decrypted_data = decrypt_symmetric_key(
            encrypted_ciphertext["ciphertext"],
            encrypted_ciphertext["nonce"],
            session_key,
        )
    except Exception as err:
        print("Symmetric decryption has failed".format(err))

    json_data = json_unpack(decrypted_data, ["msg", "status"])

    try:
        response_status = json_data["status"]
    except Exception as err:
        print("Extraing field from json data failed: {}".format(err))
        sys.exit(1)

    # If status is login succeeded then we are done...
    if DEBUG:
        print("*** sixth msg data... ***\n")
        print(json_data)
        print("\n *** check json data field... *** \n ")

    if response_status == "Login Success":
        if DEBUG:
            print("*** SUCCESS: Exiting user login... ***\n")
        return user_name, session_key, socket
    elif response_status == "Login Failure":
        if DEBUG:
            print("*** FAILURE: Exiting user login... ***\n")
        sys.exit(1)
    else:
        if DEBUG:
            print("*** UNKOWN ERROR: Exiting user login... ***\n")
        sys.exit(1)


def request_client_list(session_key, client_socket, user_name):
    """request a list of clients from server.

    The function will send the list request to server and get a list of users
    from the client along with their DH contribution.
    """
    if DEBUG:
        print("*** Request user list entry point... ***\n")

    try:
        nounce = nounce_generator()
    except Exception as err:
        print("Generating nounce failed: {}".format(err))

    json_data = json_pack(["msg", "nounce"], ["List", nounce])

    # Integrity protected encryption done here
    try:
        ciphertext, nonce = encrypt_symmetric_key(json_data, session_key)
    except Exception as err:
        print("Symmetric encryption failed: {}".format(err))

    encrypted_data = json_pack(["ciphertext", "nonce"], [ciphertext, nonce])

    list_msg = json_pack(["msg", "ciphertext", "user"],
                         ["Request", encrypted_data, user_name])
    if DEBUG:
        print("*** list_msg... ***\n")
        print(list_msg)

    try:
        client_socket.send_json(list_msg)
    except Exception as err:
        print("Sending list msg as json failed: {}".format(err))

    try:
        response = client_socket.recv_json()
    except Exception as err:
        print(
            "Recving response for the list msg as json failed: {}".format(err))

    if DEBUG:
        print("*** response... ***\n")
        print(response)

    json_data = json_unpack(response, ["msg", "ciphertext"])

    # Decrypt incoming request using the session key
    encrypted_ciphertext = json_unpack(json_data["ciphertext"],
                                       ["ciphertext", "nonce"])
    try:
        decrypted_data = decrypt_symmetric_key(
            encrypted_ciphertext["ciphertext"],
            encrypted_ciphertext["nonce"],
            session_key,
        )
    except Exception as err:
        print("Symmetric decryption has failed: {}".format(err))

    json_data = json_unpack(decrypted_data, ["user_db", "nounce_plus_one"])

    # Verify the freshness
    try:
        assert json_data["nounce_plus_one"] == nounce + 1
    except Exception as err:
        print("\nVerifying the freshness failed.".format(err))
        sys.exit(1)

    if DEBUG:
        print("*** Exiting user list request... ***\n")

    return json_data["user_db"]


def client_logout(session_key, client_socket, user_name):
    """logs out the client from the server."""
    if DEBUG:
        print("*** Entering request logout handler... ***\n")

    json_data = json_pack(["msg", "user"], ["Logout", user_name])

    # Integrity protected encryption done here
    try:
        ciphertext, nonce = encrypt_symmetric_key(json_data, session_key)
    except Exception as err:
        print("Symmetric encryption has failed.".format(err))

    encrypted_data = json_pack(["ciphertext", "nonce"], [ciphertext, nonce])

    msg = json_pack(["msg", "ciphertext", "user"],
                    ["Logout", encrypted_data, user_name])
    if DEBUG:
        print("*** logout_msg... ***\n")
        print(msg)

    try:
        client_socket.send_json(msg)
    except Exception as err:
        print("Sending msg has failed.".format(err))

    try:
        response = client_socket.recv_json()
    except Exception as err:
        print("Recving msg has failed.".format(err))

    if DEBUG:
        print("*** response... ***\n")
        print(response)

    json_data = json_unpack(response, ["msg", "ciphertext"])

    # Decrypt incoming request using the session key
    encrypted_ciphertext = json_unpack(json_data["ciphertext"],
                                       ["ciphertext", "nonce"])

    try:
        decrypted_data = decrypt_symmetric_key(
            encrypted_ciphertext["ciphertext"],
            encrypted_ciphertext["nonce"],
            session_key,
        )
    except Exception as err:
        print("Symmetric decryption has failed.".format(err))

    json_data = json_unpack(decrypted_data, ["status"])

    # Verify the response from server
    try:
        assert json_data["status"] == "Success"
    except Exception as err:
        print("\nDid not get logout success response from server.".format(err))
        sys.exit(1)

    if DEBUG:
        print("*** Exiting logout request... ***\n")


def user_send_message(user, user_db, message, my_secret, my_name):
    """Send message to another user"""
    try:
        user = user.strip().lower()
    except Exception as err:
        print("Grabbing user name failed.".format(err))

    # Arrange for user contribution to be in database
    try:
        update_list()
    except Exception as err:
        print("Grabbing user name failed.".format(err))

    if user in user_db:
        if DEBUG:
            print("\n*** user is in the database... ***\n")
        pass
    else:
        if DEBUG:
            print("\n*** user is not in the database... ***\n")
        print("User is not signed in. Try again..\n")
        return

    # Check if the user happened to change his/her contribution since last message, if any
    try:
        if user in old_contrib:
            if DEBUG:
                print("\n*** user is in the old contrib... ***\n")
            # If so, remove any information related to old contribution
            if user_db[user]["contrib"] != old_contrib[user]:
                if user in to_hashes:
                    del to_hashes[user]
        else:
            if DEBUG:
                print("\n*** user is in the old contrib... ***\n")
            pass
        old_contrib[user] = user_db[user]["contrib"]
    except Exception as err:
        print("Changing user info has failed.".format(err))

        # Get stored user contribution
    try:
        recd_user_contrib = deserialize_DH_public_key(
            user_db[user]["contrib"].encode())
    except Exception as err:
        print("Deserialize DH public key has failed.".format(err))

    # Establish session key using the user contribution
    try:
        session_key = generate_DH_shared_key(recd_user_contrib, my_secret)
    except Exception as err:
        print("Generating DH shared key has failed.".format(err))

    # Craft message
    try:
        if user in to_hashes:
            prev_hash = to_hashes[user]
        else:
            prev_hash = ""
    except Exception as err:
        print("Retrieving previous hash has failed.".format(err))

    json_data = json_pack(["msg", "prev_hash", "sender"],
                          [message, prev_hash, my_name])

    # Integrity protected encryption done here
    try:
        ciphertext, nonce = encrypt_symmetric_key(json_data, session_key)
    except Exception as err:
        print("Symmetric encryption has failed.".format(err))

    encrypted_data = json_pack(["ciphertext", "nonce"], [ciphertext, nonce])

    # Final message to be sent
    list_msg = json_pack(["msg", "ciphertext", "user"],
                         ["Chat", encrypted_data, my_name])

    # Send the message in a single TCP connection
    try:
        tcp_send_message(user_db[user]["user_ip"], user_db[user]["user_port"],
                         list_msg)
    except Exception as err:
        print("TCP sending msg has failed.".format(err))

    # Store hash of this message for crafting the next message
    try:
        to_hashes[user] = hash_message(message)
    except Exception as err:
        print("Symmetric encryption has failed.".format(err))


def print_signed_in_users():
    """Pretty format all the users currently logged in"""
    # global user_db
    print("<- Signed In Users: ", end=" ")
    for x in user_db:
        print(x, end=" ")
    print("\n")


def update_list():
    """Updates the global variable containing the database of user
    contributions etc.
    """
    global user_db, server_session_key, client_socket, user_name
    try:
        user_db = request_client_list(server_session_key, client_socket,
                                      user_name)
    except Exception as err:
        print("Requesting list of signed in clients has failed.".format(err))


def tcp_send_message(ip, port, message):
    """Send a message in a singple TCP connection"""
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.connect((ip, port))
        s.sendall(message.encode())
        s.close()
    except Exception as err:
        print("Sending msg using TCP has failed.".format(err))


def act_on_tcp_connection(conn, addr, secret):
    """Receive a single message from a TCP connection"""
    try:
        BUFFER_SIZE = 1024
        # Keep receiving until connection is closed
        recd = b""
        while 1:
            data = conn.recv(BUFFER_SIZE)
            if not data:
                break
            recd += data
        act_on_message(recd.decode(), secret)
        conn.close()
    except Exception as err:
        print("Sending msg using TCP has failed.".format(err))


def act_on_message(message, my_secret):
    """Process message received over TCP"""
    if DEBUG:
        print("*** received a message... ***\n")
        print(message)

    global user_db

    recd_msg = json_unpack(message, ["msg", "ciphertext", "user"])

    # Arrange for user contribution to be in database
    try:
        update_list()
    except Exception as err:
        print("Updating signed in users list has failed.".format(err))
        return

    # Check if the user happened to change his/her contribution since last message, if any
    try:
        if recd_msg["user"].strip().lower() in old_contrib:
            # If so, remove any information related to old contribution
            if (user_db[recd_msg["user"].strip().lower()]["contrib"] !=
                    old_contrib[recd_msg["user"].strip().lower()]):
                if recd_msg["user"].strip().lower() in from_hashes:
                    del from_hashes[recd_msg["user"].strip().lower()]
        old_contrib[recd_msg["user"].strip().lower()] = user_db[
            recd_msg["user"].strip().lower()]["contrib"]
    except Exception as err:
        print("Checking user changes has failed.".format(err))
        return

    # Get stored user contribution
    try:
        recd_user_contrib = deserialize_DH_public_key(
            user_db[recd_msg["user"].strip().lower()]["contrib"].encode())
    except Exception as err:
        print("Getting stored user contribution has failed.".format(err))
        return

    # Establish session key using the user contribution
    try:
        session_key = generate_DH_shared_key(recd_user_contrib, my_secret)
    except Exception as err:
        print("Establishing session key using DH has failed.".format(err))
        return

    if DEBUG:
        print("*** decrypting message now... ***\n")

    # Decrypt incoming chat using the session key
    encrypted_ciphertext = json_unpack(recd_msg["ciphertext"],
                                       ["ciphertext", "nonce"])

    try:
        decrypted_data = decrypt_symmetric_key(
            encrypted_ciphertext["ciphertext"],
            encrypted_ciphertext["nonce"],
            session_key,
        )
    except Exception as err:
        print("Symmetric description has failed.".format(err))
        return

    # Verify prev hash
    json_data = json_unpack(decrypted_data, ["msg", "prev_hash", "sender"])
    if DEBUG:
        print("*** verifying hash now... ***\n")

    try:
        if recd_msg["user"] in from_hashes:
            expected = from_hashes[recd_msg["user"]]
        else:
            expected = ""
        assert expected == json_data["prev_hash"]
    except Exception as err:
        print("Verifying previous hash has failed.".format(err))
        return

    # Verify the sender
    if DEBUG:
        print("*** verifying sender now... ***\n")
    try:
        assert (recd_msg["user"].strip().lower() == json_data["sender"].
                strip().lower())
    except Exception as err:
        print("Verifying the sender has failed.".format(err))
        return

    # Show the message to user
    print("\n<- <From " + recd_msg["user"] + ">: " + json_data["msg"])
    print("")

    # Store the hash for verifying next message in future
    try:
        from_hashes[recd_msg["user"]] = hash_message(json_data["msg"])
    except Exception as err:
        print("Storing verification hash has failed.".format(err))
        return

    prompt()


def listen_for_messages(secret):
    """msg listening handler.
    """

    global my_ip, my_port

    try:
        my_ip = "localhost"
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.bind((my_ip, 0))
        my_port = s.getsockname()[1]
        s.listen(10)
        while 1:
            conn, addr = s.accept()
            t = Thread(target=act_on_tcp_connection, args=(conn, addr, secret))
            t.start()
    except Exception as err:
        print("Storing verification hash has failed.".format(err))


def main():
    """main function for client.

    main() will maintain the DH session secrets and then let the user login.
    """
    # FIXME(shuwen): we need to handle the case if the server crashed but the
    # client is still trying to talk to it

    global user_db, my_ip, my_port, server_session_key, client_socket, user_name

    # arg parser
    try:
        conf = arg_parser()
        server_ip, server_port, server_pk = read_server_conf(conf)
    except Exception as err:
        print("arg parsing failed: {}".format(err))

    # create Diffie Hellman secret and contribution
    try:
        contrib, secret = DH_keys()
    except Exception as err:
        print("Main: generating DH keys has failed.".format(err))

    if DEBUG:
        print("generating secret and contrib ... \n")

    # start listening on a specific port for future messages from other users
    try:
        t = Thread(target=listen_for_messages, args=(secret, ))
        t.start()
    except Exception as err:
        print("Main: starting listener routine has failed.".format(err))

    global my_ip, my_port, user_db, server_session_key, client_socket

    while my_port == -1:
        pass

    # User login with server
    if DEBUG:
        print("*** user login... ***\n")

    try:
        user_name, server_session_key, client_socket = user_login(
            server_ip, server_port, server_pk, contrib, secret, my_ip, my_port)
    except Exception as err:
        print("Main: user login failed".format(err))
        os._exit(1)

    # Display prompt to the user
    prompt()

    try:
        while True:
            message = input()
            if DEBUG:
                print(message)
            if message == "list":
                update_list()
                print_signed_in_users()
            elif "send" in message:
                update_list()
                if len(message.split(" ")) <= 2:
                    print(
                        "You need to enter something so that we can send it.\n"
                    )
                else:
                    cmd = message.split(" ", 2)
                    user_send_message(cmd[1], user_db, cmd[2], secret,
                                      user_name)
            elif message == "logout":
                client_logout(server_session_key, client_socket, user_name)
                os._exit(1)
            else:
                print("The commands you can use are only *list* or *send*.\n")
            prompt()
    # Handle keyboard interrupt, notify server and exit from chat killing all
    # threads
    except KeyboardInterrupt:
        client_logout(server_session_key, client_socket, user_name)
        os._exit(1)


# Global variables, to facilitate access among many functions

# To store the information related to previous messages sent/recd to/from other users
to_hashes = {}
from_hashes = {}
old_contrib = {}

# Will be filled when the user will start listening for messages
my_ip = ""
my_port = -1
client_socket = None

user_name = ""

# Will be filled when user will connect to the server
user_db = {}
server_session_key = ""

if __name__ == "__main__":
    main()