#!/usr/bin/env python3
"""
Author       - pyCity
Date         - 1/22/2019
Version      - 2.0

Usage:       - python server.py --tls 127.0.0.1 4444
             - Run setup.sh to generate certs

Description: - Handler for client.py Supports TLS encryption using
             - a DHE-RSA-AES256-SHA256 cipher.
"""

import socket
import sys
import argparse
import ssl
import time


class Server:
    def __init__(self, host, port, enc):
        self.host = host
        self.port = port
        self.enc = enc
        try:
            self.s = socket.socket()
        except socket.error as hell:
            raise hell

    def create_sock(self):
        try:
            if self.enc:
                print("SSL encryption enabled")
                context = ssl.SSLContext(ssl.PROTOCOL_TLSv1_2)
                context.set_ciphers('DHE-RSA-AES256-SHA256')
                context.load_dh_params("dhparam.pem")
                context.load_cert_chain(certfile="server.crt", keyfile="server.key")
                self.s = context.wrap_socket(self.s, do_handshake_on_connect=True, server_side=True)
        except socket.error as msg:
            print("Socket creation error: {}".format(str(msg)))
            self.s.close()
            sys.exit()

    def bind_sock(self):
        attempts, max = 0, 10
        while attempts < max:
            try:
                print("Binding socket to port {}".format(str(self.port)))
                self.s.bind((self.host, self.port))
                print("Listening for clients...")
                self.s.listen(5)
                break
            except socket.error as msg:
                print("Error binding socket: {}\n  Retrying...".format(str(msg)))
                time.sleep(3)
                attempts += 1
        self.s.close()
        sys.exit()

    def handle(self):
        conn, address = self.s.accept()
        print("Connection established from: {} {}".format(address[0], str(address[1])))
        self.send_commands(conn)
        conn.close()
        self.s.close()

    def send_commands(self, conn):
        """Server loop to send commands until exit"""

        try:
            while True:
                cmd = input()
                if cmd == 'exit':
                    conn.send(str.encode("kill"))
                    break
                elif len(str.encode(cmd)) > 0:
                    conn.send(str.encode(cmd))
                    client_response = str(conn.recv(1024), "utf-8")
                    print(client_response, end="")
        except KeyboardInterrupt:
            conn.close()
            self.s.close()
            sys.exit()


if __name__ == "__main__":

    parser = argparse.ArgumentParser(description="Python tcp server")
    parser.add_argument("host", help="IP address to bind for listening")
    parser.add_argument("port", help="Port to bind for listening", type=int)
    parser.add_argument("--enc", help="Enable TLS encryption", action="store_true")
    args = parser.parse_args()
    host, port, enc = args.host, args.port, args.enc

    server = Server(host, port, enc)
    server.create_sock()
    server.bind_sock()
    server.handle()
