# -*- test-case-name: twisted.conch.test.test_filetransfer -*-
#
# Copyright (c) Twisted Matrix Laboratories.
# See LICENSE for details.


import struct, errno

from twisted.internet import defer, protocol
from twisted.python import failure, log

from common import NS, getNS
from twisted.conch.interfaces import ISFTPServer, ISFTPFile

from zope import interface



class FileTransferBase(protocol.Protocol):

    versions = (3, )

    packetTypes = {}

    def __init__(self):
        self.buf = ''
        self.otherVersion = None # this gets set

    def sendPacket(self, kind, data):
        self.transport.write(struct.pack('!LB', len(data)+1, kind) + data)

    def dataReceived(self, data):
        self.buf += data
        while len(self.buf) > 5:
            length, kind = struct.unpack('!LB', self.buf[:5])
            if len(self.buf) < 4 + length:
                return
            data, self.buf = self.buf[5:4+length], self.buf[4+length:]
            packetType = self.packetTypes.get(kind, None)
            if not packetType:
                log.msg('no packet type for', kind)
                continue
            f = getattr(self, 'packet_%s' % packetType, None)
            if not f:
                log.msg('not implemented: %s' % packetType)
                log.msg(repr(data[4:]))
                reqId, = struct.unpack('!L', data[:4])
                self._sendStatus(reqId, FX_OP_UNSUPPORTED,
                                 "don't understand %s" % packetType)
                #XXX not implemented
                continue
            try:
                f(data)
            except:
                log.err()
                continue
                reqId ,= struct.unpack('!L', data[:4])
                self._ebStatus(failure.Failure(e), reqId)

    def _parseAttributes(self, data):
        flags ,= struct.unpack('!L', data[:4])
        attrs = {}
        data = data[4:]
        if flags & FILEXFER_ATTR_SIZE == FILEXFER_ATTR_SIZE:
            size ,= struct.unpack('!Q', data[:8])
            attrs['size'] = size
            data = data[8:]
        if flags & FILEXFER_ATTR_OWNERGROUP == FILEXFER_ATTR_OWNERGROUP:
            uid, gid = struct.unpack('!2L', data[:8])
            attrs['uid'] = uid
            attrs['gid'] = gid
            data = data[8:]
        if flags & FILEXFER_ATTR_PERMISSIONS == FILEXFER_ATTR_PERMISSIONS:
            perms ,= struct.unpack('!L', data[:4])
            attrs['permissions'] = perms
            data = data[4:]
        if flags & FILEXFER_ATTR_ACMODTIME == FILEXFER_ATTR_ACMODTIME:
            atime, mtime = struct.unpack('!2L', data[:8])
            attrs['atime'] = atime
            attrs['mtime'] = mtime
            data = data[8:]
        if flags & FILEXFER_ATTR_EXTENDED == FILEXFER_ATTR_EXTENDED:
            extended_count ,= struct.unpack('!L', data[:4])
            data = data[4:]
            for i in xrange(extended_count):
                extended_type, data = getNS(data)
                extended_data, data = getNS(data)
                attrs['ext_%s' % extended_type] = extended_data
        return attrs, data

    def _packAttributes(self, attrs):
        flags = 0
        data = ''
        if 'size' in attrs:
            data += struct.pack('!Q', attrs['size'])
            flags |= FILEXFER_ATTR_SIZE
        if 'uid' in attrs and 'gid' in attrs:
            data += struct.pack('!2L', attrs['uid'], attrs['gid'])
            flags |= FILEXFER_ATTR_OWNERGROUP
        if 'permissions' in attrs:
            data += struct.pack('!L', attrs['permissions'])
            flags |= FILEXFER_ATTR_PERMISSIONS
        if 'atime' in attrs and 'mtime' in attrs:
            data += struct.pack('!2L', attrs['atime'], attrs['mtime'])
            flags |= FILEXFER_ATTR_ACMODTIME
        extended = []
        for k in attrs:
            if k.startswith('ext_'):
                ext_type = NS(k[4:])
                ext_data = NS(attrs[k])
                extended.append(ext_type+ext_data)
        if extended:
            data += struct.pack('!L', len(extended))
            data += ''.join(extended)
            flags |= FILEXFER_ATTR_EXTENDED
        return struct.pack('!L', flags) + data

class FileTransferServer(FileTransferBase):

    def __init__(self, data=None, avatar=None):
        FileTransferBase.__init__(self)
        self.client = ISFTPServer(avatar) # yay interfaces
        self.openFiles = {}
        self.openDirs = {}

    def packet_INIT(self, data):
        version ,= struct.unpack('!L', data[:4])
        self.version = min(list(self.versions) + [version])
        data = data[4:]
        ext = {}
        while data:
            ext_name, data = getNS(data)
            ext_data, data = getNS(data)
            ext[ext_name] = ext_data
        our_ext = self.client.gotVersion(version, ext)
        our_ext_data = ""
        for (k,v) in our_ext.items():
            our_ext_data += NS(k) + NS(v)
        self.sendPacket(FXP_VERSION, struct.pack('!L', self.version) + \
                                     our_ext_data)

    def packet_OPEN(self, data):
        requestId = data[:4]
        data = data[4:]
        filename, data = getNS(data)
        flags ,= struct.unpack('!L', data[:4])
        data = data[4:]
        attrs, data = self._parseAttributes(data)
        assert data == '', 'still have data in OPEN: %s' % repr(data)
        d = defer.maybeDeferred(self.client.openFile, filename, flags, attrs)
        d.addCallback(self._cbOpenFile, requestId)
        d.addErrback(self._ebStatus, requestId, "open failed")

    def _cbOpenFile(self, fileObj, requestId):
        fileId = str(hash(fileObj))
        if fileId in self.openFiles:
            raise KeyError, 'id already open'
        self.openFiles[fileId] = fileObj
        self.sendPacket(FXP_HANDLE, requestId + NS(fileId))

    def packet_CLOSE(self, data):
        requestId = data[:4]
        data = data[4:]
        handle, data = getNS(data)
        assert data == '', 'still have data in CLOSE: %s' % repr(data)
        if handle in self.openFiles:
            fileObj = self.openFiles[handle]
            d = defer.maybeDeferred(fileObj.close)
            d.addCallback(self._cbClose, handle, requestId)
            d.addErrback(self._ebStatus, requestId, "close failed")
        elif handle in self.openDirs:
            dirObj = self.openDirs[handle][0]
            d = defer.maybeDeferred(dirObj.close)
            d.addCallback(self._cbClose, handle, requestId, 1)
            d.addErrback(self._ebStatus, requestId, "close failed")
        else:
            self._ebClose(failure.Failure(KeyError()), requestId)

    def _cbClose(self, result, handle, requestId, isDir = 0):
        if isDir:
            del self.openDirs[handle]
        else:
            del self.openFiles[handle]
        self._sendStatus(requestId, FX_OK, 'file closed')

    def packet_READ(self, data):
        requestId = data[:4]
        data = data[4:]
        handle, data = getNS(data)
        (offset, length), data = struct.unpack('!QL', data[:12]), data[12:]
        assert data == '', 'still have data in READ: %s' % repr(data)
        if handle not in self.openFiles:
            self._ebRead(failure.Failure(KeyError()), requestId)
        else:
            fileObj = self.openFiles[handle]
            d = defer.maybeDeferred(fileObj.readChunk, offset, length)
            d.addCallback(self._cbRead, requestId)
            d.addErrback(self._ebStatus, requestId, "read failed")

    def _cbRead(self, result, requestId):
        if result == '': # python's read will return this for EOF
            raise EOFError()
        self.sendPacket(FXP_DATA, requestId + NS(result))

    def packet_WRITE(self, data):
        requestId = data[:4]
        data = data[4:]
        handle, data = getNS(data)
        offset, = struct.unpack('!Q', data[:8])
        data = data[8:]
        writeData, data = getNS(data)
        assert data == '', 'still have data in WRITE: %s' % repr(data)
        if handle not in self.openFiles:
            self._ebWrite(failure.Failure(KeyError()), requestId)
        else:
            fileObj = self.openFiles[handle]
            d = defer.maybeDeferred(fileObj.writeChunk, offset, writeData)
            d.addCallback(self._cbStatus, requestId, "write succeeded")
            d.addErrback(self._ebStatus, requestId, "write failed")

    def packet_REMOVE(self, data):
        requestId = data[:4]
        data = data[4:]
        filename, data = getNS(data)
        assert data == '', 'still have data in REMOVE: %s' % repr(data)
        d = defer.maybeDeferred(self.client.removeFile, filename)
        d.addCallback(self._cbStatus, requestId, "remove succeeded")
        d.addErrback(self._ebStatus, requestId, "remove failed")

    def packet_RENAME(self, data):
        requestId = data[:4]
        data = data[4:]
        oldPath, data = getNS(data)
        newPath, data = getNS(data)
        assert data == '', 'still have data in RENAME: %s' % repr(data)
        d = defer.maybeDeferred(self.client.renameFile, oldPath, newPath)
        d.addCallback(self._cbStatus, requestId, "rename succeeded")
        d.addErrback(self._ebStatus, requestId, "rename failed")

    def packet_MKDIR(self, data):
        requestId = data[:4]
        data = data[4:]
        path, data = getNS(data)
        attrs, data = self._parseAttributes(data)
        assert data == '', 'still have data in MKDIR: %s' % repr(data)
        d = defer.maybeDeferred(self.client.makeDirectory, path, attrs)
        d.addCallback(self._cbStatus, requestId, "mkdir succeeded")
        d.addErrback(self._ebStatus, requestId, "mkdir failed")

    def packet_RMDIR(self, data):
        requestId = data[:4]
        data = data[4:]
        path, data = getNS(data)
        assert data == '', 'still have data in RMDIR: %s' % repr(data)
        d = defer.maybeDeferred(self.client.removeDirectory, path)
        d.addCallback(self._cbStatus, requestId, "rmdir succeeded")
        d.addErrback(self._ebStatus, requestId, "rmdir failed")

    def packet_OPENDIR(self, data):
        requestId = data[:4]
        data = data[4:]
        path, data = getNS(data)
        assert data == '', 'still have data in OPENDIR: %s' % repr(data)
        d = defer.maybeDeferred(self.client.openDirectory, path)
        d.addCallback(self._cbOpenDirectory, requestId)
        d.addErrback(self._ebStatus, requestId, "opendir failed")

    def _cbOpenDirectory(self, dirObj, requestId):
        handle = str(hash(dirObj))
        if handle in self.openDirs:
            raise KeyError, "already opened this directory"
        self.openDirs[handle] = [dirObj, iter(dirObj)]
        self.sendPacket(FXP_HANDLE, requestId + NS(handle))

    def packet_READDIR(self, data):
        requestId = data[:4]
        data = data[4:]
        handle, data = getNS(data)
        assert data == '', 'still have data in READDIR: %s' % repr(data)
        if handle not in self.openDirs:
            self._ebStatus(failure.Failure(KeyError()), requestId)
        else:
            dirObj, dirIter = self.openDirs[handle]
            d = defer.maybeDeferred(self._scanDirectory, dirIter, [])
            d.addCallback(self._cbSendDirectory, requestId)
            d.addErrback(self._ebStatus, requestId, "scan directory failed")

    def _scanDirectory(self, dirIter, f):
        while len(f) < 250:
            try:
                info = dirIter.next()
            except StopIteration:
                if not f:
                    raise EOFError
                return f
            if isinstance(info, defer.Deferred):
                info.addCallback(self._cbScanDirectory, dirIter, f)
                return
            else:
                f.append(info)
        return f

    def _cbScanDirectory(self, result, dirIter, f):
        f.append(result)
        return self._scanDirectory(dirIter, f)

    def _cbSendDirectory(self, result, requestId):
        data = ''
        for (filename, longname, attrs) in result:
            data += NS(filename)
            data += NS(longname)
            data += self._packAttributes(attrs)
        self.sendPacket(FXP_NAME, requestId +
                        struct.pack('!L', len(result))+data)

    def packet_STAT(self, data, followLinks = 1):
        requestId = data[:4]
        data = data[4:]
        path, data = getNS(data)
        assert data == '', 'still have data in STAT/LSTAT: %s' % repr(data)
        d = defer.maybeDeferred(self.client.getAttrs, path, followLinks)
        d.addCallback(self._cbStat, requestId)
        d.addErrback(self._ebStatus, requestId, 'stat/lstat failed')

    def packet_LSTAT(self, data):
        self.packet_STAT(data, 0)

    def packet_FSTAT(self, data):
        requestId = data[:4]
        data = data[4:]
        handle, data = getNS(data)
        assert data == '', 'still have data in FSTAT: %s' % repr(data)
        if handle not in self.openFiles:
            self._ebStatus(failure.Failure(KeyError('%s not in self.openFiles'
                                        % handle)), requestId)
        else:
            fileObj = self.openFiles[handle]
            d = defer.maybeDeferred(fileObj.getAttrs)
            d.addCallback(self._cbStat, requestId)
            d.addErrback(self._ebStatus, requestId, 'fstat failed')

    def _cbStat(self, result, requestId):
        data = requestId + self._packAttributes(result)
        self.sendPacket(FXP_ATTRS, data)

    def packet_SETSTAT(self, data):
        requestId = data[:4]
        data = data[4:]
        path, data = getNS(data)
        attrs, data = self._parseAttributes(data)
        if data != '':
            log.msg('WARN: still have data in SETSTAT: %s' % repr(data))
        d = defer.maybeDeferred(self.client.setAttrs, path, attrs)
        d.addCallback(self._cbStatus, requestId, 'setstat succeeded')
        d.addErrback(self._ebStatus, requestId, 'setstat failed')

    def packet_FSETSTAT(self, data):
        requestId = data[:4]
        data = data[4:]
        handle, data = getNS(data)
        attrs, data = self._parseAttributes(data)
        assert data == '', 'still have data in FSETSTAT: %s' % repr(data)
        if handle not in self.openFiles:
            self._ebStatus(failure.Failure(KeyError()), requestId)
        else:
            fileObj = self.openFiles[handle]
            d = defer.maybeDeferred(fileObj.setAttrs, attrs)
            d.addCallback(self._cbStatus, requestId, 'fsetstat succeeded')
            d.addErrback(self._ebStatus, requestId, 'fsetstat failed')

    def packet_READLINK(self, data):
        requestId = data[:4]
        data = data[4:]
        path, data = getNS(data)
        assert data == '', 'still have data in READLINK: %s' % repr(data)
        d = defer.maybeDeferred(self.client.readLink, path)
        d.addCallback(self._cbReadLink, requestId)
        d.addErrback(self._ebStatus, requestId, 'readlink failed')

    def _cbReadLink(self, result, requestId):
        self._cbSendDirectory([(result, '', {})], requestId)

    def packet_SYMLINK(self, data):
        requestId = data[:4]
        data = data[4:]
        linkPath, data = getNS(data)
        targetPath, data = getNS(data)
        d = defer.maybeDeferred(self.client.makeLink, linkPath, targetPath)
        d.addCallback(self._cbStatus, requestId, 'symlink succeeded')
        d.addErrback(self._ebStatus, requestId, 'symlink failed')

    def packet_REALPATH(self, data):
        requestId = data[:4]
        data = data[4:]
        path, data = getNS(data)
        assert data == '', 'still have data in REALPATH: %s' % repr(data)
        d = defer.maybeDeferred(self.client.realPath, path)
        d.addCallback(self._cbReadLink, requestId) # same return format
        d.addErrback(self._ebStatus, requestId, 'realpath failed')

    def packet_EXTENDED(self, data):
        requestId = data[:4]
        data = data[4:]
        extName, extData = getNS(data)
        d = defer.maybeDeferred(self.client.extendedRequest, extName, extData)
        d.addCallback(self._cbExtended, requestId)
        d.addErrback(self._ebStatus, requestId, 'extended %s failed' % extName)

    def _cbExtended(self, data, requestId):
        self.sendPacket(FXP_EXTENDED_REPLY, requestId + data)

    def _cbStatus(self, result, requestId, msg = "request succeeded"):
        self._sendStatus(requestId, FX_OK, msg)

    def _ebStatus(self, reason, requestId, msg = "request failed"):
        code = FX_FAILURE
        message = msg
        if reason.type in (IOError, OSError):
            if reason.value.errno == errno.ENOENT: # no such file
                code = FX_NO_SUCH_FILE