"""ftpStream: Enhanced FTP with enhanced usability.

    Features:   server connect
                [FTP operations]
                open
                close
            
    $Revision: 1 $
    Released into the public domain $Date: 10/28/04 12:04p $
    Steve Holden    http://www.holdenweb.com/
"""

import ftplib

class ftpStreamError(ftplib.Error):
    pass

class InvalidModeError(ftpStreamError):
    pass

class StreamOpenError(ftpStreamError):
    pass

class ftpStream:
    "Specialised FTP with file-like interface."

    def __init__(self, host='', name='', passwd='', acct=''):
        "Create FTP object and set stream socket inactive."
        self.ftp = ftplib.FTP(host, name, passwd, acct)
        self.strsock = None

    def __getattr__(self, attrname):
        "Delegate unrecognised methods/attributes to the FTP object."
        return getattr(self.ftp, attrname)

    def open(self, file, mode="r"):
        "Open an input or output socket on a remote file."
        if self.strsock:
            raise StreamOpenError
        self.mode = mode
        if mode == "r":
            self.strsock = self.ftp.transfercmd("RETR " + file)
        elif mode == "w":
            self.strsock = self.ftp.transfercmd("STOR " + file)
        else:
            raise InvalidModeError

    def read(self, size=1024):
        "read a block from the stream socket."
        if self.mode != "r":
            raise InvalidModeError
        block = self.strsock.recv(size)
        return block

    def write(self, block):
        "write a block to the stream socket."
        if self.mode != "w":
            raise InvalidModeError
        while len(block) > 0:
            sentlen = self.strsock.send(block)
            block = block[sentlen:]

    def close(self):
        "Close stream socket but retain FTP connectivity."
        if self.strsock:
            self.strsock.close()
            self.strsock = None
            return self.ftp.voidresp() # Handle expected response
        else:
            raise ftpStreamOpenError

    def quit(self):
        "Close stream socket and quit underlying FTP object."
        if self.strsock:
            self.close()
        self.ftp.quit()

#
# Simple test if called as main routine
#
if __name__ == "__main__":
    INHOST = "127.0.0.1"
    INPATH = "/BookDir"
    INFILE = "bigfile.tgz"
    OUTHOST = "www.holdenweb.com"
    OUTPATH = "/public"
    OUTFILE = "newfile.tgz"
    INUSER = OUTUSER = "anonymous"
    INPASS = OUTPASS = "user@site"
    
    print "<<<File per stream>>>"
    # Test streams at various blocksizes
    for size in [2**x for x in range(10,16)]:
        # Establish ftp connections
        instream = ftpStream(INHOST, INUSER, INPASS)
        outstream = ftpStream(OUTHOST, OUTUSER, OUTPASS)
        instream.cwd(INPATH)
        outstream.cwd(OUTPATH)
        # Prepare sockets for transfer
        instream.open(INFILE)
        outstream.open(OUTFILE, "w")
        # Loop through transfers until done
        flen = 0
        while 1:
            block = instream.read(size)     # read next block
            if len(block) == 0:             # handle end of file
                break
            flen += len(block)              # update data length
            outstream.write(block)
        instream.quit()
        outstream.quit()
        print "Transferred", flen, "bytes in blocks of", size

    print "<<<Multiple files on one stream>>>"
    # Test multiple transfers via a single stream
    instream = ftpStream(INHOST, INUSER, INPASS)
    outstream = ftpStream(OUTHOST, OUTUSER, OUTPASS)
    instream.cwd(INPATH)
    outstream.cwd(OUTPATH)
    for size in [10**i for i in 1,2,3,4]:
        # Prepare sockets for transfer
        instream.open(INFILE)
        outstream.open(OUTFILE, "w")
        # Loop through transfers until done
        flen = 0
        while 1:
            block = instream.read(size)     # read next block
            if len(block) == 0:             # handle end of file
                break
            flen += len(block)              # update data length
            outstream.write(block)
        instream.close()
        outstream.close()
        print "Transferred", flen, "bytes in chunks of", size
    instream.quit()
    outstream.quit()