# -*- test-case-name: mv3d.test.test_upload -*-
# Copyright (C) 2010-2012 Mortal Coil Games
# See LICENSE for details.
"""
Information on Conch client usage:
http://twistedmatrix.com/documents/current/conch/howto/conch_client.html
@author:  mike
"""

import glob
import stat
import os
import time

from zope.interface import implements

from twisted.conch.client import connect, options
from twisted.conch.ssh import common, connection, filetransfer
from twisted.conch.ssh import keys, userauth
from twisted.internet import defer
from twisted.internet.defer import inlineCallbacks
from twisted.python import log
from twisted.conch.ssh import channel

from mv3d.resource.iasset import IUploader


class ClientUserAuth(userauth.SSHUserAuthClient):
    """
    Provides either password or public key authentication to the SFTP client
    """
    privateKey = None
    publicKey = None

    def __init__(self, user, instance, password=None, key=None):
        """
        key is a tuple of private, public key. They can either be strings or
        instances of Key()
        """
        userauth.SSHUserAuthClient.__init__(self, user, instance)
        self.password = password
        if key is not None:
            private, public = key
            if isinstance(private, str):
                self.privateKey = keys.Key().fromString(private)
            else:
                self.privateKey = private
            if isinstance(public, str):
                self.publicKey = keys.Key().fromString(public)
            else:
                self.publicKey = public


    def getPassword(self, _prompt=None):
        """
        If we are running with text password, then return it. Otherwise, None
        """
        if self.password is not None:
            return defer.succeed(self.password)
        #returning None will cause it to try other means


    def getPublicKey(self):
        """
        Return the public key set in init or None
        """
        return defer.succeed(self.publicKey)


    def getPrivateKey(self):
        """
        Returns the private key set in init or None
        """
        return defer.succeed(self.privateKey)



class SFTPChannel(channel.SSHChannel):
    """
    Ruthlessly copied from twisted.conch.scripts.cftp with some modifications
    """
    name = 'session'
    client = None
    closeDeferred = None

    def __init__(self, conn):
        channel.SSHChannel.__init__(self, conn=conn)
        self.dfrd = defer.Deferred()


    def channelOpen(self, _data):
        """
        Called when the connection is opened and the protocol is up.
        """
        d = self.conn.sendRequest(self, 'subsystem',
            common.NS('sftp'), wantReply=1)
        d.addCallback(self._cbSubsystem)
        d.addErrback(self.dfrd.errback)


    def _cbSubsystem(self, _result):
        self.client = filetransfer.FileTransferClient()
        self.client.makeConnection(self)
        self.dataReceived = self.client.dataReceived
        self.dfrd.callback(True)


    def closed(self):
        """
        Called when the connection is closed.
        """
        if self.closeDeferred is not None:
            self.closeDeferred.callback(None)
            self.closeDeferred = None


    def disconnect(self):
        """
        Disconnect from the server
        """
        self.closeDeferred = defer.Deferred()
        self.client.transport.loseConnection()


    def sendFile(self, local, remote="", progressCallback=None):
        """
        Send a file (or files) to the server. Wildcards may be supported in some
        form. (or they may explode the world if you use them.) Local is the
        local file name (Full path) and remote is the remote filename or path.
        """
        if '*' in local or '?' in local: # wildcard
            if remote:
                remote, _rest = self._getFilename(remote)
                path = remote
                d = self.client.getAttrs(path)
                d.addCallback(self._cbPutTargetAttrs, remote, local)
                return d
            else:
                remote = ''
                files = glob.glob(local)
                return self._cbPutMultipleNext(None, files, remote)
        if remote:
            remote, _rest = self._getFilename(remote)
        else:
            remote = os.path.split(local)[1]
        size = os.stat(local).st_size
        lf = file(local, 'rb')
        path = remote
        flags = filetransfer.FXF_WRITE | filetransfer.FXF_CREAT | filetransfer.FXF_TRUNC
        d = self.client.openFile(path, flags, {})
        d.addCallback(self._cbPutOpenFile, lf, progressCallback, size)
        return d


    def _cbPutTargetAttrs(self, attrs, path, files):
        if not stat.S_ISDIR(attrs['permissions']):
            return "Wildcard put with non-directory target."
        return self._cbPutMultipleNext(None, files, path)


    def _cbPutMultipleNext(self, _res, files, path):
        f = None
        while files and not f:
            try:
                f = files.pop(0)
                lf = file(f, 'rb')
            except:
                log.deferr()
                f = None
        if not f:
            return
        name = os.path.split(f)[1]
        remote = "/".join([path, name])
        log.msg((name, remote, path))
        flags = filetransfer.FXF_WRITE | filetransfer.FXF_CREAT | filetransfer.FXF_TRUNC
        d = self.client.openFile(remote, flags, {})
        d.addCallback(self._cbPutOpenFile, lf)
        #d.addErrback(self._ebCloseLf, lf)
        d.addBoth(self._cbPutMultipleNext, files, path)
        return d


    def _cbPutOpenFile(self, rf, lf, progressCallback=None, size=None):
        numRequests = 2
        dList = []
        chunks = []
        startTime = time.time()
        for chunk in range(numRequests):
            d = self._cbPutWrite(None, rf, lf, chunks, startTime, progressCallback, size)
            if d:
                dList.append(d)
        dl = defer.DeferredList(dList, fireOnOneErrback=1)
        dl.addCallback(self._cbPutDone, rf, lf)
        return dl


    def _cbPutWrite(self, _ignored, rf, lf, chunks, startTime, progressCallback=None, size=None):
        chunk = self._getNextChunk(chunks)
        start, size = chunk
        lf.seek(start)
        if progressCallback:
            progressCallback(start / float(size))
#        arr = array("c")
#        try:
#            arr.fromfile(lf, size)
#        except EOFError:
#            pass
#        data = arr.tostring()
        data = lf.read(size)
        if data:
            d = rf.writeChunk(start, data)
            if d is None:
                return
            d.addCallback(self._cbPutWrite, rf, lf, chunks, startTime)
            return d
        else:
            return


    def _cbPutDone(self, _ignored, rf, lf):
        lf.close()
        rf.close()
        return 'Transferred %s to %s' % (lf.name, rf.name)


    def _getNextChunk(self, chunks):
        end = 0
        for chunk in chunks:
            if end == 'eof':
                return # nothing more to get
            if end != chunk[0]:
                i = chunks.index(chunk)
                chunks.insert(i, (end, chunk[0]))
                return (end, chunk[0] - end)
            end = chunk[1]
        bufSize = 32768
        chunks.append((end, end + bufSize))
        return (end, bufSize)


    def _getFilename(self, line):
        line.lstrip()
        if not line:
            return None, ''
        if line[0] in '\'"':
            ret = []
            line = list(line)
            try:
                for i in range(1, len(line)):
                    c = line[i]
                    if c == line[0]:
                        return ''.join(ret), ''.join(line[i + 1:]).lstrip()
                    elif c == '\\': # quoted character
                        del line[i]
                        if line[i] not in '\'"\\':
                            raise IndexError, "bad quote: \\%s" % line[i]
                        ret.append(line[i])
                    else:
                        ret.append(line[i])
            except IndexError:
                raise IndexError, "unterminated quote"
        ret = line.split(None, 1)
        if len(ret) == 1:
            return ret[0], ''
        else:
            return ret



class ClientConnection(connection.SSHConnection):
    options = None
    channel = None

    def serviceStarted(self):
        """
        Called when the connection is made, creates a new channel which can be
        accessed by our channel attribute.
        """
        self.channel = SFTPChannel(conn=self)
        self.openChannel(self.channel)



class SFTPClient(object):
    """
    A simple and easy SFTP client implementation.

    Example usage:
    >>> sftp = SFTPClient("some.host.com", "mike", password="password")
    >>> d = sftp.sendFile("/path/to/some/file.txt")
    >>> d.addCallback(lambda _: sftp.disconnect())

    """
    connection = None

    def __init__(self, host, user, password=None, key=None, port=22):
        self.host = host
        self.user = user
        self.password = password
        self.key = key
        self.port = port


    @inlineCallbacks
    def connect(self):
        """
        Will connect automatically when first used, but this method can be
        called manually to force a connection (or reconnection). 
        """
        def verifyHostKey(_transport, _host, _pubKey, _fingerprint):
            """ we should check host keys eventually """
            return defer.succeed(1)
        conn = ClientConnection()
        uao = ClientUserAuth(self.user, conn, self.password, self.key)
        opts = options.ConchOptions()
        opts["host"] = self.host
        opts["port"] = self.port
        opts["user-authentications"] = ["password", "publickey"]
        conn.options = opts
        yield connect.connect(self.host, self.port, opts, verifyHostKey, uao)
        self.connection = conn.channel
        yield conn.channel.dfrd


    @inlineCallbacks
    def sendFile(self, local, remote=None, progressCallback=None):
        """
        Send a file to the remote server
        """
        if self.connection is None:
            yield self.connect()
        try:
            yield self.connection.sendFile(local, remote, progressCallback)
        except UnicodeDecodeError:
            self.disconnect()
            raise


    def disconnect(self):
        """
        Disconnect from the server
        """
        connection = self.connection
        self.connection = None
        if connection is not None:
            return connection.disconnect()



class SFTPUploader(object):
    """
    Asset data uploader that uses SFTP
    """
    implements(IUploader)

    client = None
    baseDir = None
    downloadUrl = None


    def __str__(self):
        return self.name

    def configure(self, cfg, section):
        """
        Configure based on the section in a config file
        """
        self.name = section
        self.host = cfg.get(section, "host")
        self.user = cfg.get(section, "user")
        self.password = None
        self.key = None
        self.port = 22
        if cfg.has_option(section, "password"):
            self.password = cfg.get(section, "password")
        if cfg.has_option(section, "publicKey"):
            pubKey = keys.Key.fromFile(cfg.get(section, "publicKey"))
            privKey = keys.Key.fromFile(cfg.get(section, "privateKey"))
            self.key = (privKey, pubKey)
        self.client = SFTPClient(self.host, self.user, self.password, self.key,
            self.port)
        if cfg.has_option(section, "baseDir"):
            self.baseDir = cfg.get(section, "baseDir")
        else:
            self.baseDir = "."
        if cfg.has_option(section, "downloadUrl"):
            self.downloadUrl = cfg.get(section, "downloadUrl")


    def upload(self, localFile, remoteFile, progressCallback=None):
        """
        Upload the file and return a deferred that will fire when the upload
        is finished
        """
        return self.client.sendFile(localFile, "/".join([self.baseDir,
            remoteFile]), progressCallback)


    def disconnect(self):
        """
        Disconnect the uploader
        """
        self.client.disconnect()