#!/usr/local/bin/python
# FileTransfer.py - CGI file transfer, may require Python 1.5 or higher
# Author: JeffBauer@bigfoot.com
# Prior art:  Aaron Watters, Jim Fulton

import os, sys, traceback
import httplib
import tempfile
from whrandom import randint
from string import joinfields
from time import time

class FileUploadAcquisition:
    """
    Server-side portion of the CGI file transfer.  For security
    reasons uploaded file should not be stored by the user-specified
    filename.
    """
    def __init__(self, environ=os.environ):
        print "Content-type: text/html"
        print "Expires: Monday, 1-Jan-96 00:00:00 GMT"
        print "Pragma: no-cache"
        print
        sys.stderr = sys.stdout
        self.tmpfile = None
        try:
            self.process()
            self.disposition()
            if self.tmpfile:
                os.unlink(self.tmpfile)
        except:
            print "<pre>"
            traceback.print_exc()
    def disposition(self):
        """
        Overwrite this method to determine what to do after the
        upload, because self.tmpfile will be removed after this
        method is invoked.
        """
        pass
    def process(self):
        import cgi
        fs = cgi.FieldStorage()
        uf = fs['file']
        # The next 3 are assigned as instance variables so the
        # disposition() method can use them.
        self.filename = fs.filename  # the user's specified filename
        self.headers = uf.headers
        self.tmpfile = tempfile.mktemp()
        f = open(self.tmpfile, 'wb')
        while 1:
            line = uf.file.readline()
            if line:
                f.write(line)
            else:
                break
        f.close()

class FileUploadRequest:
    """
    Client-side CGI file upload.  Forces binary mode for a single
    file.  Could be extended to handle multiple files, if desirable.
    """
    def __init__(self, uri, host, port):
        self.uri = uri
        self.host = host
        self.port = port
        self.queryString = None
        self.boundary= '%s%s_%s_%s' % \
                        ('-----', int(time()), os.getpid(), randint(1,10000))
    def load(self, filename, headerDict=None):
        f = open(filename, 'rb')
        data = f.read()
        f.close()
        hdr = []; part = []; total = []
        hdr.append('Content-Disposition: form-data; name="file"; filename="%s"' % (filename))
        hdr.append('Content-Type: application/octet-stream')
        hdr.append('Content-Length: %d' % len(data))
        # Add user-defined header values here, if provided.
        # Reminder: FieldStorage will convert key to lowercase.
        if type(headerDict) == type({}):
            for k in headerDict.keys():
                hdr.append("%s: %s", (k, headerDict[k]))
        part.append("%s\n\n%s" % (joinfields(hdr,'\n'), data))
        total.append('--%s\n' % self.boundary)
        total.append(joinfields(part, "\n--%s\n" % self.boundary))
        total.append('\n--%s--\n' % self.boundary)
        self.queryString = joinfields(total, '')
    def request(self):
        query = self.queryString
        contentType = 'multipart/form-data; boundary=%s' % self.boundary
        contentLength = str(len(query))
        h = httplib.HTTP()
        h.connect(self.host, self.port)
        h.putrequest('POST', self.uri)
        h.putheader('Accept', '*/*')
        h.putheader('Proxy-Connection', 'Keep-Alive')
        h.putheader('User-Agent', 'Bond/007 [en] (WinNT; U)')
        h.putheader('Content-Type', contentType)
        h.putheader('Content-Length', contentLength)
        h.endheaders()
        h.send(query)
        rcode, rmsg, headers = h.getreply()
        response = h.getfile().read()
        if rcode != 200:
            msg = "error: %s, %s\n%s %s" % (rcode, self.uri, rmsg, response)
            raise FileUploadRequestException(msg)
        return response

def test():
    usage = 'usage: cgiupload filename'
    if len(sys.argv) < 2:
        print usage
        sys.exit(1)
    filename = sys.argv[1]

    uri = '/cgi-bin/cgiupload.py'
    host = 'localhost'
    port = 80

    fileUploadRequest = FileUploadRequest(uri, host, port)
    fileUploadRequest.load(filename)
    response = fileUploadRequest.request()
    print response

if __name__ == '__main__':
    if os.environ.has_key('SERVER_PROTOCOL'):
        FileUploadAcquisition(os.environ)
    else:
        test()