#!/usr/bin/env python3

import sys
import argparse
import socket
import ssl
import operator

# Set up and parse CLI arguments
def arg_setup():
    parser = argparse.ArgumentParser()
    parser.add_argument('-p', type=int, default=27998,
        help='TCP port to connect to. Defaults to 27998.')
    parser.add_argument('-s', action='store_true',
        help='Flag to require secure SSL connections.')
    parser.add_argument('hostname', help='Hostname.')
    parser.add_argument('neu_id', help='NEU ID.')
    
    return parser.parse_args()

# Create the socket using the given port and hostname. Make the socket secure
#   if indicated
def socket_setup(port, hostname, secure):
    host = socket.gethostbyname(hostname)
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    
    if(secure):
        try:
            wrapped = ssl.wrap_socket(s)
            wrapped.connect((host, port))
            
            return wrapped
        except:
            sys.exit("Failed to create SSL connection at host {} and port {}."
                .format(hostname, port))
    else:
        try:
            s.connect((host, port))
        
            return s
        except:
            sys.exit("Failed to create connection at host {} and port {}."
                .format(hostname, port))

if __name__ == "__main__":
    # Parse CLI arguments
    args = arg_setup()
    # Secure connections use port 27999 by default
    if(args.s & (args.p == 27998)):
        args.p = 27999
    
    # Create socket to connect to appropriate hostname and port
    s = socket_setup(args.p, args.hostname, args.s)
    
    # 1. Send HELLO message
    s.send(bytes("cs3700spring2018 HELLO {}\n".format(args.neu_id), 'utf-8'))
    
    # 2. Receive first message and split the input (by spaces) to parse into
    #      the math problem
    data = s.recv(256)
    splitdata = data.split()
    
    # Map string representations to actual math operators, including proper
    #   floor division
    ops = {"+": operator.add,
        "-": operator.sub,
        "*": operator.mul,
        "/": operator.floordiv}
    
    # 3. Continue processing math problems until a message is received that is
    #      not a STATUS message
    while((splitdata[0] == b"cs3700spring2018") & (splitdata[1] == b"STATUS")):
        # Convert byte operator to str, and byte numbers to ints
        op = splitdata[3].decode()
        num1 = int(splitdata[2])
        num2 = int(splitdata[4])
        
        # Get the answer to the math problem
        ans = ops[op](num1, num2)
        
        # 4. Send answer to the server
        s.send(bytes("cs3700spring2018 {}\n".format(ans), 'utf-8'))
        
        # 5. Receive next message and split to parse
        data = s.recv(256)
        splitdata = data.split()

    # 6. Check that the message is the BYE message, and output the secret flag
    if(splitdata[2] == b"BYE"):
        print(splitdata[1].decode())
        
        s.close()
        
        sys.exit(0)
    
    # 7. If a message is received in an incorrect format, error:
    s.close()
    sys.exit("Received incorrect message. Ending communication. Message:\n{}".format(data))