__author__ = 'JackSong'

import paramiko
import logging,os
import chardet,time
from ioloop import IOloop
from tornado.options import options
from tornado.websocket import WebSocketClosedError
from paramiko.ssh_exception import AuthenticationException,SSHException,BadHostKeyException
from paramiko import RSAKey


class Bridge(object):
    def __init__(self,websocket):
        self._websocket = websocket
        self._shell = None
        self._fd = 0
        self.ssh = paramiko.SSHClient()
        self._status = False

    def open_file(self,ip,port,username):
        time_str = time.strftime("%Y-%m-%d-%H-%M-%S",time.localtime())
        return open(r'%s\\%s-%s-%s-%s.log' % (options.log_dir,ip,port,username,time_str),"a+")

    def write_file(self,message):
        me = self

        def run():
            if message != None:
                res = me._file.write(message)
                me._file.flush()
                return res

        return run

    def close_file(self):
        if hasattr(self,"_file"):
            self._file.close()

    @property
    def fd(self):
        return self._fd

    @property
    def status(self):
        return self._status

    @property
    def websocket(self):
        return self._websocket

    @property
    def shell(self):
        return self._shell

    def open(self,data={}):
        from utils.tools import PrpcryptInitData
        self.ssh.load_system_host_keys()
        self.ssh.set_missing_host_key_policy(
            paramiko.AutoAddPolicy())
        data = PrpcryptInitData(**data.to_dict())
        self.connect(data,data.key_file)

    def connect(self,data,key_filename=None):
        try:
            key = None
            if key_filename:
                key_file=os.path.join(options.file_path,key_filename)
                key = RSAKey.from_private_key_file(key_file)
            self.ssh.connect(
                hostname = data.ip,
                port = int(data.port),
                username = data.username,
                password = data.passwd,
                timeout = 3,
                gss_auth = False,
                gss_kex = False,
                gss_deleg_creds = False,
                pkey = key
            )
            self._file = self.open_file(data.ip,data.port,data.username)
        except AuthenticationException:
            raise Exception("auth failed user:%s ,passwd:%s" %
                            (data.username,data.passwd))
        except BadHostKeyException:
            raise Exception('Unable to establish a link:%s',data.hostname)
        except SSHException:
            raise Exception("could not connect to host:%s:%s" %
                            (data.username,data.passwd))
        self.establish()

    def establish(self,term="xterm"):
        self._shell = self.ssh.invoke_shell(term)
        self._shell.setblocking(0)
        self._fd = self._shell.fileno()
        self._status = True
        IOloop.current().register(self)

    def trans_forward(self,data=""):
        if self._shell:
            try:
                self._shell.send(data)
                return True
            except:
                return False
        else:
            return False

    def trans_back(self,result):
        if self._websocket:
            try:
                if type(result) == bytes:
                    detect = chardet.detect(result)
                    encoding = detect.get('encoding')
                    if result and encoding and encoding != 'utf-8':
                        result = result.decode(encoding).encode("utf-8")
                    self._websocket.write_message(result)
                    if result.strip() == 'logout':
                        return False
                    return True
                if type(result) == str:
                    self._websocket.write_message(result)
                    if result.strip() == 'logout':
                        return False
                    return True
            except WebSocketClosedError as e:
                logging.debug(e)
                return False

    def destroy(self):
        self._websocket.close()
        self.ssh.close()
        self.close_file()
        IOloop.current().remove_handler(self.fd)